solver.py

import os
import torch.nn as nn
import torch.optim as optim
from torch.autograd import Variable
from torchvision import models
from torch.utils.data import DataLoader
import utils
#from logger import Logger
from model import *

from data_loader import *
from torchvision import transforms
import numpy as np

class Solver(object):
    def __init__(self, args):
        # parameters
        self.model_name = args.model_name
        self.patch_size = args.patch_size
        self.num_threads = args.num_threads
        self.exposure_value = args.exposure_value
        self.num_channels = args.num_channels

        self.num_epochs = args.num_epochs
        self.save_epochs = args.save_epochs
        self.batch_size = args.batch_size
        self.test_batch_size = args.test_batch_size
        self.lr = args.lr

        self.train_dataset = args.train_dataset
        self.test_dataset = args.test_dataset
        
        self.save_dir = args.save_dir
        self.gpu_mode = args.gpu_mode

        self.stride = args.stride
     
        self.build_model()

    def build_model(self):
        # networks
        self.stopup_G = Generator(num_channels=self.num_channels, base_filter=64, stop='up')
        self.stopdown_G = Generator(num_channels=self.num_channels, base_filter=64, stop='down')
        
        self.stopup_D = NLayerDiscriminator(num_channels=2*self.num_channels,base_filter=64, image_size=self.patch_size)
        self.stopdown_D = NLayerDiscriminator(num_channels=2*self.num_channels,base_filter=64, image_size=self.patch_size)

        print('---------- Networks architecture -------------')
        utils.print_network(self.stopup_G)
        utils.print_network(self.stopdown_D)
        print('----------------------------------------------')

        # weigh initialization
        self.stopup_G.weight_init()
        self.stopdown_G.weight_init()
 
        self.stopup_D.weight_init()
        self.stopdown_D.weight_init()        

        # optimizer
        self.stopup_G_optimizer = optim.Adam(self.stopup_G.parameters(), lr=self.lr, betas=(0.5, 0.999))
        self.stopdown_G_optimizer = optim.Adam(self.stopdown_G.parameters(), lr=self.lr, betas=(0.5, 0.999)) 

        self.stopup_D_optimizer = optim.Adam(self.stopup_D.parameters(), lr=self.lr, betas=(0.5, 0.999))
        self.stopdown_D_optimizer = optim.Adam(self.stopdown_D.parameters(), lr=self.lr, betas=(0.5, 0.999)) 

        # loss function
        if self.gpu_mode:
            self.stopup_G = nn.DataParallel(self.stopup_G)
            self.stopdown_G = nn.DataParallel(self.stopdown_G)
            
            self.stopup_D = nn.DataParallel(self.stopup_D)
            self.stopdown_D = nn.DataParallel(self.stopdown_D)

            self.stopup_G.cuda()
            self.stopdown_G.cuda()
 
            self.stopup_D.cuda()
            self.stopdown_D.cuda()

            self.L1_loss = nn.L1Loss().cuda()
            self.criterionGAN = GANLoss().cuda()
        
        else:
            self.L1_loss = nn.L1Loss()
            self.MSE_loss = nn.MSELoss()
            self.BCE_loss = nn.BCELoss()
            self.criterionGAN = GANLoss()

        return

    def load_dataset(self, dataset, is_train=True):
        if self.num_channels == 1:
            is_gray = True
        else:
            is_gray = False

        if is_train:
            print('Loading train datasets...')
            train_set = get_loader(self.train_dataset)
            return DataLoader(dataset=train_set, num_workers=self.num_threads, batch_size=self.batch_size,
                              shuffle=True)
        else:
            print('Loading test datasets...')
            test_set = get_loader(self.test_dataset)
            return DataLoader(dataset=test_set, num_workers=self.num_threads,
                              batch_size=self.test_batch_size,
                              shuffle=False)

    def train(self):
        # load dataset
        train_data_loader = self.load_dataset(dataset=self.train_dataset, is_train=True)
        test_data_loader = self.load_dataset(dataset=self.test_dataset[0], is_train=False)

        # set the logger
        stopup_G_log_dir = os.path.join(self.save_dir, 'stopup_G_logs')
        if not os.path.exists(stopup_G_log_dir):
            os.mkdir(stopup_G_log_dir)
        #stopup_G_logger = Logger(stopup_G_log_dir)

        stopup_D_log_dir = os.path.join(self.save_dir, 'stopup_D_logs')
        if not os.path.exists(stopup_D_log_dir):
            os.mkdir(stopup_D_log_dir)
        #stopup_D_logger = Logger(stopup_D_log_dir)


        stopdown_G_log_dir = os.path.join(self.save_dir, 'stopdown_G_logs')
        if not os.path.exists(stopdown_G_log_dir):
            os.mkdir(stopdown_G_log_dir)
        #stopdown_G_logger = Logger(stopdown_G_log_dir)

        stopdown_D_log_dir = os.path.join(self.save_dir, 'stopdown_D_logs')
        if not os.path.exists(stopdown_D_log_dir):
            os.mkdir(stopdown_D_log_dir)
        #stopdown_D_logger = Logger(stopdown_D_log_dir)


        ################# Pre-train generator #################
        self.epoch_pretrain = 10

        # Load pre-trained parameters of generator
        if not self.load_model(is_pretrain=True):
            # Pre-training generator for 10 epochs
            print('Pre-training is started.')
            self.stopup_G.train()
            self.stopdown_G.train()

            for epoch in range(self.epoch_pretrain):
                for iter, (lr, hr) in enumerate(train_data_loader):
                    # input data (low dynamic image)
                    if self.num_channels == 1:
                        x_ = Variable(utils.norm(hr[:, 0].unsqueeze(1), vgg=False))
                        y_ = Variable(utils.norm(lr[:, 0].unsqueeze(1), vgg=False))
                    else:
                        x_ = Variable(utils.norm(hr, vgg=False))
                        y_ = Variable(utils.norm(lr, vgg=False))

                    if self.gpu_mode:
                        x_ = x_.cuda()
                        y_ = y_.cuda()

                    # Train generator
                    self.stopup_G_optimizer.zero_grad()
                    self.stopdown_G_optimizer.zero_grad()                    

                    '''
                    stopup
                    '''
                    stopup_est = self.stopup_G(y_)

                    # Content losses
                    stopup_content_loss = self.L1_loss(stopup_est, x_)
                    stopup_G_loss = stopup_content_loss

                    stopup_G_loss.backward()
                    self.stopup_G_optimizer.step()

                    '''
                    stopdown
                    '''
                    stopdown_est = self.stopdown_G(x_)

                    # Content losses
                    stopdown_content_loss = self.L1_loss(stopdown_est, y_)
                    stopdown_G_loss = stopdown_content_loss

                    stopdown_G_loss.backward()
                    self.stopdown_G_optimizer.step()

                    # log
                    print("Epoch: [%2d] [%4d/%4d] stopup_G: %.6f/stopdown_G: %.6f"
                          % ((epoch + 1), (iter + 1), len(train_data_loader), stopup_G_loss.data[0], stopdown_G_loss.data[0]), end='\r')

                    # siyeong
                    if (iter % 100 == 0):
                        import random
                        index = random.randrange(0,self.batch_size)
                         
                        input_data = torch.cat((y_[index], x_[index]), 1)
                        est_data = torch.cat((stopup_est[index], stopdown_est[index]),1)
                        square = torch.cat((input_data, est_data), 2)                   
                        square = utils.denorm(square.cpu().data, vgg=False)

                        square_img = transforms.ToPILImage()(square)

                        square_img.show()

            print('Pre-training is finished.')

            # Save pre-trained parameters of generator
            self.save_model(is_pretrain=True)

        ################# Adversarial train #################
        print('Training is started.')
        # Avg. losses
        stopup_G_avg_loss = []
        stopup_D_avg_loss = []
        stopdown_G_avg_loss = []
        stopdown_D_avg_loss = []
 
        step = 0

        # test image
        test_lr, test_hr = test_data_loader.dataset.__getitem__(2)
        test_lr = test_lr.unsqueeze(0)
        test_hr = test_hr.unsqueeze(0)

        self.stopup_G.train()
        self.stopup_D.train()

        self.stopdown_G.train()
        self.stopdown_D.train()

        for epoch in range(self.num_epochs):
            # learning rate is decayed by a factor of 10 every 20 epoch
            if (epoch + 1) % 20 == 0:
                for param_group in self.stopup_G_optimizer.param_groups:
                    param_group["lr"] /= 2.0
                print("Learning rate decay for G: lr={}".format(self.stopup_G_optimizer.param_groups[0]["lr"]))
                for param_group in self.stopup_D_optimizer.param_groups:
                    param_group["lr"] /= 2.0
                print("Learning rate decay for D: lr={}".format(self.stopup_D_optimizer.param_groups[0]["lr"]))

                for param_group in self.stopdown_G_optimizer.param_groups:
                    param_group["lr"] /= 2.0
                print("Learning rate decay for G: lr={}".format(self.stopdown_G_optimizer.param_groups[0]["lr"]))
                for param_group in self.stopdown_D_optimizer.param_groups:
                    param_group["lr"] /= 2.0
                print("Learning rate decay for D: lr={}".format(self.stopdown_D_optimizer.param_groups[0]["lr"]))

            stopup_G_epoch_loss = 0
            stopup_D_epoch_loss = 0

            stopdown_G_epoch_loss = 0
            stopdown_D_epoch_loss = 0

            for iter, (lr, hr) in enumerate(train_data_loader):
                # input data (low dynamic image)
                mini_batch = lr.size()[0]

                if self.num_channels == 1:
                    x_ = Variable(utils.norm(hr[:, 0].unsqueeze(1), vgg=False))
                    y_ = Variable(utils.norm(lr[:, 0].unsqueeze(1), vgg=False))
                else:
                    x_ = Variable(utils.norm(hr, vgg=False))
                    y_ = Variable(utils.norm(lr, vgg=False))

                if self.gpu_mode:
                    x_ = x_.cuda()
                    y_ = y_.cuda()
                    # labels
                    real_label = Variable(torch.ones(mini_batch).cuda())
                    fake_label = Variable(torch.zeros(mini_batch).cuda())
                else:
                    # labels
                    real_label = Variable(torch.ones(mini_batch))
                    fake_label = Variable(torch.zeros(mini_batch))

                # Reset gradient
                self.stopup_D_optimizer.zero_grad()
                self.stopdown_D_optimizer.zero_grad()

                # Train discriminator with real data
                stopup_D_real_decision = self.stopup_D(torch.cat((x_, y_),1))
                stopdown_D_real_decision = self.stopdown_D(torch.cat((y_, x_),1))

                stopup_D_real_loss = self.criterionGAN(stopup_D_real_decision, True)
                stopdown_D_real_loss = self.criterionGAN(stopdown_D_real_decision, True)

                # Train discriminator with fake data
                stopup_est = self.stopup_G(y_)
                stopdown_est = self.stopdown_G(x_)

                stopup_D_fake_decision = self.stopup_D(torch.cat((stopup_est, y_),1))
                stopdown_D_fake_decision = self.stopdown_D(torch.cat((stopdown_est, x_),1))

                stopup_D_fake_loss = self.criterionGAN(stopup_D_fake_decision, False)
                stopdown_D_fake_loss = self.criterionGAN(stopdown_D_fake_decision, False)

                stopup_D_loss = 0.5*stopup_D_real_loss + 0.5*stopup_D_fake_loss
                stopdown_D_loss = 0.5*stopdown_D_real_loss + 0.5*stopdown_D_fake_loss

                # Back propagation
                stopup_D_loss.backward(retain_graph=True)
                stopdown_D_loss.backward(retain_graph=True)

                self.stopup_D_optimizer.step()
                self.stopdown_D_optimizer.step()

                # Reset gradient
                self.stopup_G_optimizer.zero_grad()
                self.stopdown_G_optimizer.zero_grad()

                # Train generator
                stopup_est = self.stopup_G(y_)
                stopdown_est = self.stopdown_G(x_)

                stopup_D_fake_decision = self.stopup_D(torch.cat((stopup_est, y_), 1))
                stopdown_D_fake_decision = self.stopdown_D(torch.cat((stopdown_est, x_), 1))

                # Adversarial loss
                stopup_GAN_loss = self.criterionGAN(stopup_D_fake_decision, True)
                stopdown_GAN_loss = self.criterionGAN(stopdown_D_fake_decision, True)

                # Content losses
                stopup_mae_loss = self.L1_loss(stopup_est, x_) 
                stopdown_mae_loss = self.L1_loss(stopdown_est, y_)

                # Total loss 
                stopup_G_loss =  stopup_mae_loss + 1e-2*stopup_GAN_loss 
                stopdown_G_loss =  stopdown_mae_loss + 1e-2*stopdown_GAN_loss 

                stopup_G_loss.backward()
                self.stopup_G_optimizer.step()

                stopdown_G_loss.backward()
                self.stopdown_G_optimizer.step()

                # siyeong
                if (iter % 100 == 0):
                    import random
                    index = random.randrange(0,self.batch_size)
                         
                    input_data = torch.cat((y_[index], x_[index]), 1)
                    est_data = torch.cat((stopup_est[index], stopdown_est[index]),1)

                    square = torch.cat((input_data, est_data), 2)

                    square = utils.denorm(square.cpu().data, vgg=False)
                    square_img = transforms.ToPILImage()(square)

                    square_img.show()

                # log
                stopup_G_epoch_loss += stopup_G_loss.data[0]
                stopup_D_epoch_loss += stopup_D_loss.data[0]

                stopdown_G_epoch_loss += stopdown_G_loss.data[0]
                stopdown_D_epoch_loss += stopdown_D_loss.data[0]

                print("Epoch: [%02d] [%05d/%05d] stopup_G/D: %.6f/%.6f, stopdown_G/D: %.6f/%.6f"
                      % ((epoch + 1), (iter + 1), len(train_data_loader), stopup_G_loss.data[0], stopup_D_loss.data[0], stopdown_G_loss.data[0], stopdown_D_loss.data[0]), end="\r")

                # tensorboard logging
                stopup_G_logger.scalar_summary('losses', stopup_G_loss.data[0], step + 1)
                stopup_D_logger.scalar_summary('losses', stopup_D_loss.data[0], step + 1)

                stopdown_G_logger.scalar_summary('losses', stopdown_G_loss.data[0], step + 1)
                stopdown_D_logger.scalar_summary('losses', stopdown_D_loss.data[0], step + 1)
 
                step += 1

            # avg. loss per epoch
            stopup_G_avg_loss.append(stopup_G_epoch_loss / len(train_data_loader))
            stopup_D_avg_loss.append(stopup_D_epoch_loss / len(train_data_loader))

            stopdown_G_avg_loss.append(stopdown_G_epoch_loss / len(train_data_loader))
            stopdown_D_avg_loss.append(stopdown_D_epoch_loss / len(train_data_loader))

            self.save_model(epoch + 1)

        # Plot avg. loss
        utils.plot_loss([stopup_G_avg_loss, stopup_D_avg_loss, stopdown_G_avg_loss, stopdown_D_avg_loss], self.num_epochs, save_dir=self.save_dir)
        print("Training is finished.")

        # Save final trained parameters of model
        self.save_model(epoch=None)

    # siyeong3
    def test(self, input_path='./', out_path='./Result/', extend = 3):
        # load model
        self.load_model(is_pretrain=False)
        scenes = listdir(input_path)
        for i, scene in enumerate(scenes):
            scene_path = join(input_path, scene)
            if not os.path.isdir(out_path):
                os.mkdir(out_path)
            
            out_name = os.path.splitext(os.path.split(scene_path)[1])[0]
            storage_path = out_path + out_name + '/'

            # mkdir storage folder
            if not os.path.isdir(storage_path):
                os.mkdir(storage_path)

            # cp middle exposure file
            cmd = "cp " + scene_path + " " + storage_path + out_name + '_EV0.png'            
            os.system(cmd)

            target = scene_path
            for i in range(1, extend+1):
                reconst = self.image_single(target, True)
                output_name = storage_path + out_name + '_EV%d' %i + '.png'
                reconst.save(output_name)

                target = output_name

            target = scene_path
            for i in range(1, extend+1):
                reconst = self.image_single(target, False)
                output_name = storage_path + out_name +'_EV-%d' %i + '.png'
                reconst.save(output_name)

                target = output_name
            print('\tImage [', out_name, '] is finished.')
        print('Test is finishied.')

    def image_single(self, img_fn, stopup):
        # load data
        img = Image.open(img_fn).convert('RGB')
        
        img = img.resize((256, 256), 4)
        tensor = transforms.ToTensor()(img)
        tensor_norm = Variable(utils.norm(tensor, vgg=False))
        tensor_expand = tensor_norm.unsqueeze(0)
        
        if stopup:
            self.stopup_G.train()
            recon_norm = self.stopup_G(tensor_expand)

        else:
            self.stopdown_G.train()
            recon_norm =  self.stopdown_G(tensor_expand)       

        recon = utils.denorm(recon_norm.cpu().data, vgg=False)
        recon = recon.squeeze(0)
        recon = torch.clamp(recon, min=0, max=1)
        recon_img = transforms.ToPILImage()(recon)
        return recon_img

    def save_model(self, epoch=None, is_pretrain=False):
        model_dir = os.path.join(self.save_dir, 'model')
        if not os.path.exists(model_dir):
            os.mkdir(model_dir)

        if is_pretrain:
            torch.save(self.stopup_G.state_dict(), model_dir + '/' + self.model_name + '_stopup_G_param_pretrain.pkl')
            torch.save(self.stopdown_G.state_dict(), model_dir + '/' + self.model_name + '_stopdown_G_param_pretrain.pkl')
 
            print('Pre-trained generator model is saved.')
        else:
            if epoch is not None:
                torch.save(self.stopup_G.state_dict(), model_dir + '/' + self.model_name +
                           '_stopup_G_param_ch%d_batch%d_epoch%d_lr%.g.pkl'
                           % (self.num_channels, self.batch_size, epoch, self.lr))
                torch.save(self.stopup_D.state_dict(), model_dir + '/' + self.model_name +
                           '_stopup_D_param_ch%d_batch%d_epoch%d_lr%.g.pkl'
                           % (self.num_channels, self.batch_size, epoch, self.lr))

                torch.save(self.stopdown_G.state_dict(), model_dir + '/' + self.model_name +
                           '_stopdown_G_param_ch%d_batch%d_epoch%d_lr%.g.pkl'
                           % (self.num_channels, self.batch_size, epoch, self.lr))
                torch.save(self.stopdown_D.state_dict(), model_dir + '/' + self.model_name +
                           '_stopdown_D_param_ch%d_batch%d_epoch%d_lr%.g.pkl'
                           % (self.num_channels, self.batch_size, epoch, self.lr))

            else:
                torch.save(self.stopup_G.state_dict(), model_dir + '/' + self.model_name +
                           '_stopup_G_param_ch%d_batch%d_epoch%d_lr%.g.pkl'
                           % (self.num_channels, self.batch_size, self.num_epochs, self.lr))
                torch.save(self.stopup_D.state_dict(), model_dir + '/' + self.model_name +
                           '_stopup_D_param_ch%d_batch%d_epoch%d_lr%.g.pkl'
                           % (self.num_channels, self.batch_size, self.num_epochs, self.lr))
                torch.save(self.stopdown_G.state_dict(), model_dir + '/' + self.model_name +
                           '_stopdown_G_param_ch%d_batch%d_epoch%d_lr%.g.pkl'
                           % (self.num_channels, self.batch_size, self.num_epochs, self.lr))
                torch.save(self.stopdown_D.state_dict(), model_dir + '/' + self.model_name +
                           '_stopdown_D_param_ch%d_batch%d_epoch%d_lr%.g.pkl'
                           % (self.num_channels, self.batch_size, self.num_epochs, self.lr))
 
            print('Trained models are saved.')

    def load_model(self, is_pretrain=False):
        model_dir = os.path.join(self.save_dir, 'model')

        if is_pretrain:
            flag_stopup = False
            flag_stopdown = False

            model_name_stopup = model_dir + '/' + self.model_name + '_stopup_G_param_pretrain.pkl'
            model_name_stopup_D = model_dir + '/' + self.model_name + '_stopup_D_param_pretrain.pkl'
 
            if os.path.exists(model_name_stopup):
                self.stopup_G.load_state_dict(torch.load(model_name_stopup))
                self.stopup_D.load_state_dict(torch.load(model_name_stopup_D))
                flag_stopup = True

            model_name_stopdown = model_dir + '/' + self.model_name + '_stopdown_G_param_pretrain.pkl'
            model_name_stopdown_D = model_dir + '/' + self.model_name + '_stopdown_D_param_pretrain.pkl'
 
            if os.path.exists(model_name_stopdown):
                self.stopdown_G.load_state_dict(torch.load(model_name_stopdown))
                self.stopdown_D.load_state_dict(torch.load(model_name_stopdown_D))

                flag_stopdown = True

            print ("[loding] (up):", flag_stopup, ', (down):',flag_stopdown)
            print (model_name_stopup)
            print (model_name_stopdown)

            if flag_stopdown and flag_stopup:
                print('Pre-trained generator model is loaded.')
                return True
            else:
                return False

        else:
            flag_stopup = False
            flag_stopdown = False

            model_name_stopup = model_dir + '/' + self.model_name + \
                         '_stopup_G_param_ch%d_batch%d_epoch%d_lr%.g.pkl' \
                         % (self.num_channels, self.batch_size, self.num_epochs, self.lr)
            print(model_name_stopup)

            if os.path.exists(model_name_stopup):
                self.stopup_G.load_state_dict(torch.load(model_name_stopup))
                flag_stopup = True

            model_name_stopdown = model_dir + '/' + self.model_name + \
                         '_stopdown_G_param_ch%d_batch%d_epoch%d_lr%.g.pkl' \
                         % (self.num_channels, self.batch_size, self.num_epochs, self.lr)
 
            if os.path.exists(model_name_stopdown):
                self.stopdown_G.load_state_dict(torch.load(model_name_stopdown))
                flag_stopdown = True
                
            print ("[loding] (up):", flag_stopup, ', (down):',flag_stopdown)
            print (model_name_stopup)
            print (model_name_stopdown)

            if flag_stopup and flag_stopdown:
               print('Trained generator model is loaded.')
               return True

            else:
               return False