train.py

import os
import glob
import socket
import logging

import tensorflow as tf
import neuralgymGit.neuralgym as ng

from inpaint_model import InpaintCAModel
from inpaint_ops import data_augument


logger = logging.getLogger()


def multigpu_graph_def(model, data, config, gpu_id=0, loss_type='g',batch_mask=None):
    with tf.device('/cpu:0'):
        images = data
    if gpu_id == 0 and loss_type == 'g':
        _, _, losses = model.build_graph_with_losses(
            images, config, summary=True, reuse=True,batch_mask=batch_mask)
    else:
        _, _, losses = model.build_graph_with_losses(
            images, config, reuse=True,batch_mask=batch_mask)
    if loss_type == 'g':
        return losses['g_loss']
    elif loss_type == 'd':
        return losses['d_loss']
    else:
        raise ValueError('loss type is not supported.')


if __name__ == "__main__":
    
    masks=None
    config = ng.Config('inpaint.yml')
    if config.GPU_ID != -1:
        ng.set_gpus(config.GPU_ID)
    else:
        ng.get_gpus(config.NUM_GPUS)

    # training data    
    if config.CUSTOM_MASK:
        # Read training image paths
        with open(config.DATA_FLIST[config.DATASET][0]) as f:
            fnames = f.read().splitlines()
            
        # Read training masks paths    
        with open(config.DATA_FLIST[config.DATASET][2]) as f1:
            fnames_mask = f1.read().splitlines()
            
        fnames_combined= [(fnames[k],fnames_mask[k]) for k in range(len(fnames))]
        data = ng.data.DataFromFNames(
            fnames_combined, [config.IMG_SHAPES,[config.IMG_SHAPES[0],config.IMG_SHAPES[0],3]], random_crop=config.RANDOM_CROP,
            fn_preprocess=None)
        
        images,masks = data.data_pipeline(config.BATCH_SIZE)
        print('masks.shape',masks)

    else:
        with open(config.DATA_FLIST[config.DATASET][0]) as f:
            fnames = f.read().splitlines()
        data = ng.data.DataFromFNames(
            fnames, config.IMG_SHAPES, random_crop=config.RANDOM_CROP,
            fn_preprocess=data_augument)
        images = data.data_pipeline(config.BATCH_SIZE)
                
    # main model
    model = InpaintCAModel()
    g_vars, d_vars, losses = model.build_graph_with_losses(
        images,batch_mask=masks, config=config)

    # validation images
    if config.VAL:
        with open(config.DATA_FLIST[config.DATASET][1]) as f:
            val_fnames = f.read().splitlines()
        with open(config.DATA_FLIST[config.DATASET][3]) as f:
            val_mask_fnames = f.read().splitlines()
        # progress monitor by visualizing static images
        for i in range(config.STATIC_VIEW_SIZE):
            static_fnames = val_fnames[i:i+1]
            static_mask_fnames = val_mask_fnames[i:i+1]
            
            static_images = ng.data.DataFromFNames(
                static_fnames, config.IMG_SHAPES, nthreads=1,
                random_crop=config.RANDOM_CROP,fn_preprocess=None).data_pipeline(1)
            
            static_masks = ng.data.DataFromFNames(
                static_mask_fnames, [config.IMG_SHAPES[0],config.IMG_SHAPES[1],3], nthreads=1,
                random_crop=config.RANDOM_CROP,fn_preprocess=None).data_pipeline(1)
            
            static_inpainted_images = model.build_static_infer_graph(
                static_images, config, name='static_view/%d' % i,mask=static_masks)

    # training settings
    lr = tf.get_variable(
        'lr', shape=[], trainable=False,
        initializer=tf.constant_initializer(1e-4))
    d_optimizer = tf.train.AdamOptimizer(lr, beta1=0.5, beta2=0.9)
    g_optimizer = d_optimizer

    # gradient processor
    if config.GRADIENT_CLIP:
        gradient_processor = lambda grad_var: (
            tf.clip_by_average_norm(grad_var[0], config.GRADIENT_CLIP_VALUE),
            grad_var[1])
    else:
        gradient_processor = None

    # log dir
    log_prefix = 'model_logs/' + '_'.join([
        ng.date_uid(), socket.gethostname(), config.DATASET,
        'MASKED' if config.GAN_WITH_MASK else 'NORMAL',
        config.GAN, config.LOG_DIR])

    # train discriminator with secondary trainer, should initialize before
    # primary trainer.
    discriminator_training_callback = ng.callbacks.SecondaryTrainer(
        pstep=1,
        optimizer=d_optimizer,
        var_list=d_vars,
        max_iters=1,
        graph_def=multigpu_graph_def,
        graph_def_kwargs={
            'model': model, 'data': images, 'config': config, 'loss_type': 'd','batch_mask':masks},
    )

    # train generator with primary trainer
    trainer = ng.train.Trainer(
        optimizer=g_optimizer,
        var_list=g_vars,
        max_iters=config.MAX_ITERS,
        graph_def=multigpu_graph_def,
        grads_summary=config.GRADS_SUMMARY,
        gradient_processor=gradient_processor,
        graph_def_kwargs={
            'model': model, 'data': images, 'config': config, 'loss_type': 'g','batch_mask':masks},
        spe=config.TRAIN_SPE,
        log_dir=log_prefix,
    )

    # add all callbacks
    if not config.PRETRAIN_COARSE_NETWORK:
        trainer.add_callbacks(discriminator_training_callback)
    trainer.add_callbacks([
        ng.callbacks.WeightsViewer(),
        ng.callbacks.ModelRestorer(trainer.context['saver'], dump_prefix='model_logs/'+config.MODEL_RESTORE+'/snap', optimistic=True),
        ng.callbacks.ModelSaver(config.TRAIN_SPE, trainer.context['saver'], log_prefix+'/snap'),
        ng.callbacks.SummaryWriter((config.VAL_PSTEPS//1), trainer.context['summary_writer'], tf.summary.merge_all()),
    ])

    # launch training
    trainer.train()