train.py

import argparse
import time
import datetime
from data import CustomDataLoader
from data.super_dataset import SuperDataset

from models import create_model
from configs import parse_config
from utils.util import print_losses, check_path, make_grid, AverageMeter
from utils.data_utils import check_old_config_val_possible
import os
import torch.multiprocessing as mp
import torch.distributed as dist
from torch.utils.tensorboard import SummaryWriter
import copy
import sys



def main():
    # parse arguments
    parser = argparse.ArgumentParser(description='Style Master')
    parser.add_argument('--cfg_file', type=str, default='./exp/sp2pII-phase4.yaml')
    parser.add_argument('--image_prompt', type=str, default='./example/reference/01.png')
    parser.add_argument('--text_prompt', type=str, default='sketch with black pencil')
    parser.add_argument('--name', type=str, default='01')
    
    args = parser.parse_args()

    # parse config
    config = parse_config(args.cfg_file)

    if args.image_prompt:
        config['training']['image_prompt'] = args.image_prompt
    if args.text_prompt:
        config['training']['text_prompt'] = args.text_prompt
    if args.name:
        config['common']['name'] = args.name

    for group in config:
        print(group + ':')
        for k, v in config[group].items():
            print('  {}: {}'.format(k, v))

    # we want GPU ids match nvidia-smi output order, so do some manipulations here.
    # GPU ids need to always start from 0, but the system variable CUDA_VISIBLE_DEVICES can be set to e.g. GPU 2 and 3.
    gpu_string = ','.join(map(str, config['common']['gpu_ids']))
    gpu_ids_fix = list(range(len(config['common']['gpu_ids'])))
    os.environ['CUDA_VISIBLE_DEVICES'] = gpu_string
    config['common']['gpu_ids'] = gpu_ids_fix

    if config['training']['DDP']:
        num_nodes = config['training']['num_nodes']
        num_gpus = len(config['common']['gpu_ids'])
        config['training']['world_size'] = num_gpus * num_nodes
        os.environ['MASTER_ADDR'] = config['training']['DDP_address']
        os.environ['MASTER_PORT'] = config['training']['DDP_port']
        mp.spawn(train_val, nprocs=num_gpus, args=(config,))
    else:
        # under DP mode, shall set batch size to (actual effective batch size * num_gpu)
        config['dataset']['batch_size'] *= len(config['common']['gpu_ids'])
        if 'patch_batch_size' in config['dataset']:
            config['dataset']['patch_batch_size'] *= len(config['common']['gpu_ids'])
        train_val(None, config)

# GPU parameter is automatically filled when using DDP. It is an irrelevant placeholder if not using DDP.
def train_val(gpu, config):
    import torch
    if config['training']['val']:
        config_val = copy.deepcopy(config)
        config_val['common']['phase'] = 'val'

    if config['training']['DDP']:
        dist.init_process_group(
            backend='nccl',
            init_method='env://',
            world_size=config['training']['world_size'],
            rank=gpu
        )
        torch.cuda.set_device(gpu)  # DDP requirement

    # Dataset and dataloader construction.
    # prepare train data
    train_dataset = SuperDataset(config, shuffle=True, check_all_data=config['dataset']['check_all_data'],
                                 DDP_device=gpu)
    if config['dataset']['train_data'] == {}:
        print("old data config format detected. Converting to new format.")
        train_dataset.config = train_dataset.convert_old_config_to_new()

    train_dataset.static_data.load_static_data()
    train_dataset.static_data.create_transforms()
    if train_dataset.check_all_data:
        train_dataset.check_data()

    if len(train_dataset) == 0:
        if gpu == 0 or gpu is None:
            print("Train set has 0 data samples. Exiting.")
        sys.exit(0)

    _, train_video_dataset = train_dataset.split_data('count', 5, mode='copy_partial')

    # prepare val data
    force_use_train_data = False
    if config['training']['val']:
        val_dataset = SuperDataset(config_val)
        if len(config_val['dataset']['train_data']) == 0:
            if check_old_config_val_possible(config_val):
                val_dataset.convert_old_config_to_new()
                val_dataset.static_data.load_static_data()
            else:
                force_use_train_data = True
        else:
            val_dataset.static_data.load_static_data()


        if len(val_dataset) == 0 or force_use_train_data:
            if gpu == 0 or gpu is None:
                print("Validation set has 0 data samples. Using part of training data for validation.")
            validation_ratio = config['training']['val_percent']/100

            train_dataset, val_dataset = train_dataset.split_data('ratio', validation_ratio, mode='split')

            if len(val_dataset) == 0:
                print("There are too few training data to establish a validation set. "
                      "Use the training set as validation set.")
                val_dataset = copy.deepcopy(train_dataset)

        val_dataset.static_data.create_transforms()
        if val_dataset.check_all_data:
            val_dataset.check_data()

    # print dataset info
    if gpu == 0 or gpu is None:
        print("--------train dataset static data content-----------")
        for i, cnt in enumerate(train_dataset.static_data.len_of_groups):
            print("%s:   %d" % (train_dataset.static_data.group_names[i], cnt))
        print("----------------------------------------")

        if config['training']['val']:
            print("--------val dataset static data content-----------")
            for i, cnt in enumerate(val_dataset.static_data.len_of_groups):
                print("group %s:   %d" % (val_dataset.static_data.group_names[i], cnt))
            print("----------------------------------------")


    # prepare dataloaders
    train_dataloader = CustomDataLoader(config, train_dataset, DDP_gpu=gpu, drop_last=config['dataset']['drop_last'])
    if config['training']['val']:
        val_dataloader = CustomDataLoader(config_val, val_dataset, DDP_gpu=gpu, drop_last=config['dataset']['drop_last'])
    config_train_video = copy.deepcopy(config)
    config_train_video['dataset']['serial_batches'] = True
    config_train_video['batch_size'] = 1
    train_video_dataloader = CustomDataLoader(config_train_video, train_video_dataset, DDP_gpu=None, drop_last=False)

    if config['training']['DDP']:
        model = create_model(config, DDP_device=gpu)      # create a DDP model given opt.model and other options
        model.setup(config, DDP_device=gpu)  # regular setup: load and print networks; create schedulers
    else:
        model = create_model(config)  # create a model (singleGPU or dataparallel) given opt.model and other options
        model.setup(config)               # regular setup: load and print networks; create schedulers
    total_iters = model.total_iters                # the total number of training iterations

    # visualization setups
    if gpu == 0 or gpu is None:
        log_dir = os.path.join(config['training']['log_dir'], config['common']['name'] + '_' +
                               datetime.datetime.now().strftime("%Y_%m_%d__%H_%M_%S"))
        check_path(log_dir)
        writer = SummaryWriter(log_dir)
    train_loss_averages = {}
    val_loss_averages = {}

    if config['training']['epoch_as_iter']:
        iter_ub = config['training']['n_epochs'] + config['training']['n_epochs_decay']
        while total_iters < iter_ub:
            # Training loop
            if gpu == 0 or gpu is None:
                print('-------------------------------Training----------------------------------')
            model.train()
            quit = False
            for i, data in enumerate(train_dataloader):
                total_iters += 1
                if total_iters >= iter_ub:
                    quit = True
                    break
                model.curr_epoch = total_iters
                model.total_iters = total_iters

                model.set_input(data)
                model.optimize_parameters()

                losses = model.get_current_losses()
                for k, v in losses.items():
                    if k not in train_loss_averages:
                        train_loss_averages[k] = AverageMeter()
                    train_loss_averages[k].update(v)

                if (total_iters % config['training']['print_freq'] == 0) and ((gpu == 0) or (gpu is None)):
                    print_losses(total_iters, 1, 1, losses, train_loss_averages)
                    tmp = make_grid(model)
                    for j in range(0, len(tmp[0])):
                        img_grid, name = tmp[0][j], tmp[1][j]
                        if config['training']['use_new_log']:
                            writer.add_image('Training/' + name, img_grid, total_iters)
                        else:
                            writer.add_image('iteration ' + str(total_iters) + ' training ' + name, img_grid)
                    for k, v in losses.items():
                        writer.add_scalar('Training/' + k, train_loss_averages[k].avg(), total_iters)

                if (total_iters % config['training']['save_latest_freq'] == 0) and ((gpu == 0) or (gpu is None)):
                    print('saving the latest model (total_iters %d)' % total_iters)
                    model.save_networks('latest')

                if (total_iters % config['training']['save_epoch_freq'] == 0) and ((gpu == 0) or (gpu is None)):
                    print('saving the model at the end of iters %d' % total_iters)
                    model.save_networks('latest')
                    #model.save_networks(total_iters)

                model.update_learning_rate()
                for k, v in losses.items():
                    train_loss_averages[k].clear()

            if quit:
                exit(0)

            if config['training']['val']:
                # Validation loop
                if gpu == 0 or gpu is None:
                    print('-------------------------------Validating----------------------------------')
                model.eval()
                for i, data in enumerate(val_dataloader):
                    with torch.no_grad():
                        model.set_input(data)
                        model.eval_step()

                        losses = model.get_current_losses()
                        for k, v in losses.items():
                            if k not in val_loss_averages:
                                val_loss_averages[k] = AverageMeter()
                            val_loss_averages[k].update(v)

                if gpu == 0 or gpu is None:
                    print_losses(total_iters, 1, 1, losses, val_loss_averages, mode='validating')
                    tmp = make_grid(model)
                    for j in range(0, len(tmp[0])):
                        img_grid, name = tmp[0][j], tmp[1][j]
                        if config['training']['use_new_log']:
                            writer.add_image('Validation/' + name, img_grid, total_iters)
                        else:
                            writer.add_image('iteration ' + str(total_iters) + ' validating ' + name, img_grid)
                    for k, v in losses.items():
                        writer.add_scalar('Validation/' + k, val_loss_averages[k].avg(), total_iters)

                main_loss = 'G' if 'G' in losses else model.loss_names[0]
                if val_loss_averages[main_loss].avg() < model.best_val_loss and (gpu == 0 or gpu is None):
                    model.best_val_loss = val_loss_averages[main_loss].avg()
                    print('New validation best loss. saving the model.')
                    model.save_networks('', val_loss=model.best_val_loss)

                for k, v in losses.items():
                    val_loss_averages[k].clear()

            if config['training']['save_training_progress']:
                # produce images on the same images every epoch to visualize how training is progressing.
                if gpu == 0 or gpu is None:
                    for i, data in enumerate(train_video_dataloader):
                        with torch.no_grad():
                            model.set_input(data)
                            model.forward()
                            tmp = make_grid(model)
                            for j in range(0, len(tmp[0])):
                                img_grid, name = tmp[0][j], tmp[1][j]
                                if config['training']['use_new_log']:
                                    writer.add_image('Training Video/' + name + ' ' + str(i), img_grid, total_iters)
                                else:
                                    writer.add_image('epoch 0 iteration ' + str(total_iters) + ' training_video ' + name + ' ' + str(i), img_grid)


    for epoch in range(model.curr_epoch, config['training']['n_epochs'] + config['training']['n_epochs_decay'] + 1):    # outer loop for different epochs; we save the model by <epoch_count>, <epoch_count>+<save_latest_freq>
        epoch_start_time = time.time()  # timer for entire epoch
        iter_data_time = time.time()    # timer for data loading per iteration
        epoch_iter = 0                  # the number of training iterations in current epoch, reset to 0 every epoch
        model.curr_epoch = epoch

        # Training loop
        if gpu == 0 or gpu is None:
            print('-------------------------------Training----------------------------------')
        model.train()
        for i, data in enumerate(train_dataloader):  # inner loop within one epoch
            iter_start_time = time.time()  # timer for computation per iteration

            total_iters += 1
            epoch_iter += 1
            model.set_input(data)         # unpack data from dataset and apply preprocessing
            model.optimize_parameters()   # calculate loss functions, get gradients, update network weights

            losses = model.get_current_losses()
            for k, v in losses.items():
                if k not in train_loss_averages:
                    train_loss_averages[k] = AverageMeter()
                train_loss_averages[k].update(v)

            if epoch_iter % config['training']['print_freq'] == 0 and ((gpu == 0) or (gpu is None)):
                print_losses(epoch, epoch_iter, (len(train_dataset) + config['dataset']['batch_size'] - 1) // config['dataset']['batch_size'], losses, train_loss_averages)
                tmp = make_grid(model)
                for j in range(0, len(tmp[0])):
                    img_grid, name = tmp[0][j], tmp[1][j]
                    if config['training']['use_new_log']:
                        writer.add_image('Training/' + name, img_grid, total_iters)
                    else:
                        writer.add_image('epoch ' + str(epoch) + ' iteration ' + str(total_iters) + ' training ' + name, img_grid)
                for k, v in losses.items():
                    writer.add_scalar('Training/' + k, train_loss_averages[k].avg(), total_iters)

            if total_iters % config['training']['save_latest_freq'] == 0 and ((gpu == 0) or (gpu is None)):   # cache our latest model every <save_latest_freq> iterations
                print('saving the latest model (epoch %d, total_iters %d)' % (epoch, total_iters))
                model.save_networks('latest')

            iter_data_time = time.time()

        if gpu == 0 or gpu is None:
            print_losses(epoch, epoch_iter, (len(train_dataset) + config['dataset']['batch_size'] - 1) // config['dataset']['batch_size'], losses, train_loss_averages)
            print('End of epoch %d / %d \t Time Taken: %d sec' % (epoch, config['training']['n_epochs'] + config['training']['n_epochs_decay'], time.time() - epoch_start_time))
            tmp = make_grid(model)
            for j in range(0, len(tmp[0])):
                img_grid, name = tmp[0][j], tmp[1][j]
                if config['training']['use_new_log']:
                    writer.add_image('Training/' + name, img_grid, total_iters)
                else:
                    writer.add_image('epoch ' + str(epoch) + ' iteration ' + str(total_iters) + ' training ' + name, img_grid)
            for k, v in losses.items():
                writer.add_scalar('Training/' + k, train_loss_averages[k].avg(), total_iters)

        if epoch % config['training']['save_epoch_freq'] == 0 and ((gpu == 0) or (gpu is None)):              # cache our model every <save_epoch_freq> epochs
            print('saving the model at the end of epoch %d, iters %d' % (epoch, total_iters))
            model.save_networks('latest')
            model.save_networks(epoch)

        model.update_learning_rate()  # update learning rates at the end of every epoch.
        for k, v in losses.items():
            train_loss_averages[k].clear()

        if config['training']['val']:
            # Validation loop
            if gpu == 0 or gpu is None:
                print('-------------------------------Validating----------------------------------')
            model.eval()
            epoch_iter = 0
            for i, data in enumerate(val_dataloader):  # inner loop within one epoch
                epoch_iter += 1

                with torch.no_grad():
                    model.set_input(data)
                    model.eval_step()

                    losses = model.get_current_losses()
                    for k, v in losses.items():
                        if k not in val_loss_averages:
                            val_loss_averages[k] = AverageMeter()
                        val_loss_averages[k].update(v)

            if gpu == 0 or gpu is None:
                print_losses(epoch, epoch_iter, (len(val_dataset) + config['dataset']['batch_size'] - 1)
                             // config['dataset']['batch_size'], losses, val_loss_averages, mode='validating')
                tmp = make_grid(model)
                for j in range(0, len(tmp[0])):
                    img_grid, name = tmp[0][j], tmp[1][j]
                    if config['training']['use_new_log']:
                        writer.add_image('Validation/' + name, img_grid, total_iters)
                    else:
                        writer.add_image('epoch ' + str(epoch) + ' iteration ' + str(total_iters) + ' validating ' + name, img_grid)
                for k, v in losses.items():
                    writer.add_scalar('Validation/' + k, val_loss_averages[k].avg(), total_iters)

            main_loss = 'G' if 'G' in losses else model.loss_names[0]
            if val_loss_averages[main_loss].avg() < model.best_val_loss and gpu == 0 or gpu is None:
                model.best_val_loss = val_loss_averages[main_loss].avg()
                print('New validation best loss. saving the model.')
                model.save_networks('', val_loss=model.best_val_loss)

            for k, v in losses.items():
                val_loss_averages[k].clear()

        if config['training']['save_training_progress']:
            # produce images on the same images every epoch to visualize how training is progressing.
            if gpu == 0 or gpu is None:
                for i, data in enumerate(train_video_dataloader):
                    with torch.no_grad():
                        model.set_input(data)
                        model.forward()
                        tmp = make_grid(model)
                        for j in range(0, len(tmp[0])):
                            img_grid, name = tmp[0][j], tmp[1][j]
                            if config['training']['use_new_log']:
                                writer.add_image('Training Video/' + name + ' ' + str(i), img_grid, total_iters)
                            else:
                                writer.add_image('epoch ' + str(epoch) + ' iteration ' + str(total_iters) + ' training_video ' + name + ' ' + str(i), img_grid)

    # If we shut down process now, writer could save incomplete data. Wait a bit to let it finish.
    time.sleep(5)

if __name__ == '__main__':
    main()