# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # BicycleGAN
#
# ## 1. Import Libs

# +
import torch
import torch.nn as nn
from torch.autograd import Variable
import torch.optim as optim
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader

from dataloader import Edges2handbags
import util
import os
import argparse

from model import *

CUDA = torch.cuda.is_available()
# -

# ## 2. Settings

# +
parser = argparse.ArgumentParser([])

parser.add_argument('--root', required=False, default='./dataset/edges2shoes/edges2shoes')
parser.add_argument('--result_dir', default='result/edges2shoes')
parser.add_argument('--weight_dir', default='weight/edges2shoes')
parser.add_argument('--load_weight', type=bool, default=False)
parser.add_argument('--batch_size', type=int, default=2)
parser.add_argument('--test_size', type=int, default=20)
parser.add_argument('--test_img_num', type=int, default=5)
parser.add_argument('--img_size', type=int, default=128)
parser.add_argument('--num_epoch', type=int, default=100)
parser.add_argument('--save_every', type=int, default=1000)
parser.add_argument('--lr', type=float, default=0.0002)
parser.add_argument('--beta_1', type=float, default=0.5)
parser.add_argument('--beta_2', type=float, default=0.999)
parser.add_argument('--lambda_kl', type=float, default=0.01)
parser.add_argument('--lambda_img', type=int, default=10)
parser.add_argument('--lambda_z', type=float, default=0.5)
parser.add_argument('--z_dim', type=int, default=8)

params = parser.parse_args([])
print(params)


# -

# ## 3. Utils

# +
def mse_loss(score, target=1):
    dtype = type(score)
    
    if target == 1:
        label = util.var(torch.ones(score.size()), requires_grad=False)
    elif target == 0:
        label = util.var(torch.zeros(score.size()), requires_grad=False)
    
    criterion = nn.MSELoss().cuda()
    loss = criterion(score, label)
    
    return loss

def L1_loss(pred, target):
    return torch.mean(torch.abs(pred - target))

def lr_decay_rule(epoch, start_decay=100, lr_decay=100):
    decay_rate = 1.0 - (max(0, epoch - start_decay) / float(lr_decay))
    return decay_rate


def save_weight(epoch=None):
    if epoch is None:
        d_cVAE_name = 'D_cVAE.pkl'
        d_cLR_name = 'D_cLR.pkl'
        g_name = 'G.pkl'
        e_name = 'E.pkl'
    else:
        d_cVAE_name = '{epochs}-{name}'.format(epochs=str(epoch), name='D_cVAE.pkl')
        d_cLR_name = '{epochs}-{name}'.format(epochs=str(epoch), name='D_cLR.pkl')
        g_name = '{epochs}-{name}'.format(epochs=str(epoch), name='G.pkl')
        e_name = '{epochs}-{name}'.format(epochs=str(epoch), name='E.pkl')
            
    torch.save(D_cVAE.state_dict(), os.path.join(self.weight_dir, d_cVAE_name))
    torch.save(D_cVAE.state_dict(), os.path.join(self.weight_dir, d_cLR_name))
    torch.save(G.state_dict(), os.path.join(self.weight_dir, g_name))
    torch.save(E.state_dict(), os.path.join(self.weight_dir, e_name))


# -

# ## 4. Load Dataset
# ### 4.1 Preprocessing

transform = transforms.Compose([
    transforms.Scale((params.img_size, params.img_size)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])

# ### 4.2 DataLoader

# +
train_dataset = Edges2handbags(params.root, transform=transform, mode='train')
test_dataset = Edges2handbags(params.root, transform=transform, mode='val')

train_data_loader = DataLoader(train_dataset, batch_size=params.batch_size, shuffle=True, num_workers=4, drop_last=True)
test_data_loader = DataLoader(test_dataset, batch_size=params.batch_size, shuffle=True, num_workers=4, drop_last=True)

train_sample = train_dataset.__getitem__(11)
print(train_sample)
print(len(train_sample))
# -

# ## 5. Build Models & Optimizers & Loss Function

# ### 5.1 Models

# +
D_cVAE = Discriminator().cuda()
D_cLR = Discriminator().cuda()
G = Generator(z_dim=params.z_dim).cuda()
E = Encoder(z_dim=params.z_dim).cuda()

print(D_cVAE)
print(D_cLR)
print(G)
print(E)
# -

# ### 5.2 Optimizers

optim_D_cVAE = optim.Adam(D_cVAE.parameters(), lr=params.lr, betas=(params.beta_1, params.beta_2))
optim_D_cLR = optim.Adam(D_cLR.parameters(), lr=params.lr, betas=(params.beta_1, params.beta_2))
optim_G = optim.Adam(G.parameters(), lr=params.lr, betas=(params.beta_1, params.beta_2))
optim_E = optim.Adam(E.parameters(), lr=params.lr, betas=(params.beta_1, params.beta_2))

# 20 x 5 x 8
fixed_z = util.var(torch.randn(params.test_size, params.test_img_num, params.z_dim)).cuda()

# ## 6. Load Pretrained Models

if params.load_weight is True:
    D_cVAE.load_state_dict(torch.load(os.path.join(params.weight_dir, 'D_cVAE.pkl')))
    D_cLR.load_state_dict(torch.load(os.path.join(params.weight_dir, 'D_cLR.pkl')))
    G.load_state_dict(torch.load(os.path.join(params.weight_dir, 'G.pkl')))
    E.load_state_dict(torch.load(os.path.join(params.weight_dir, 'E.pkl')))

D_cVAE.train()
D_cLR.train()
G.train()
E.train()

# +
D_A_avg_losses = []
G_GAN_avg_losses = []
KL_div_avg_losses = []
img_recon_avg_losses = []
G_alone_avg_losses = []

for epoch in range(1, params.num_epoch + 1):
    D_A_losses = []
    G_GAN_losses = []
    KL_div_losses = []
    img_recon_losses = []
    G_alone_losses = []
    for iters, (img, ground_truth) in enumerate(train_data_loader): # 1
        # img : (2, 3, 128, 128) of domain A / ground_truth : (2, 3, 128, 128) of domain B
        img, ground_truth = util.var(img), util.var(ground_truth)
        
        # Seperate data for cVAE_GAN and cLR_GAN
        cVAE_data = {'img' : img[0].unsqueeze(dim=0), 'ground_truth' : ground_truth[0].unsqueeze(dim=0)}
        cLR_data = {'img' : img[1].unsqueeze(dim=0), 'ground_truth' : ground_truth[1].unsqueeze(dim=0)}
        
        ''' ----------------------------- 1. Train D ----------------------------- '''
        #############   Step 1. D loss in cVAE-GAN #############
        
        # Encoded latent vector
        mu, log_variance = E(cVAE_data['ground_truth'])
        std = torch.exp(log_variance / 2)
        random_z = util.var(torch.randn(1, params.z_dim))
        encoded_z = (random_z * std) + mu
        
        # Generate fake image
        fake_img_cVAE = G(cVAE_data['img'], encoded_z)
        
        # Get scores and loss
        real_d_cVAE_1, real_d_cVAE_2 = D_cVAE(cVAE_data['ground_truth'])
        fake_d_cVAE_1, fake_d_cVAE_2 = D_cVAE(fake_img_cVAE)
                
        # mse_loss for LSGAN
        D_loss_cVAE_1 = mse_loss(real_d_cVAE_1, 1) + mse_loss(fake_d_cVAE_1, 0)
        D_loss_cVAE_2 = mse_loss(real_d_cVAE_2, 1) + mse_loss(fake_d_cVAE_2, 0)
                
        #############   Step 2. D loss in cLR-GAN   #############

        # Random latent vector
        random_z = util.var(torch.randn(1, params.z_dim))

        # Generate fake image
        fake_img_cLR = G(cLR_data['img'], random_z)

        # Get scores and loss
        real_d_cLR_1, real_d_cLR_2 = D_cLR(cLR_data['ground_truth'])
        fake_d_cLR_1, fake_d_cLR_2 = D_cLR(fake_img_cLR)
        
        D_loss_cLR_1 = mse_loss(real_d_cLR_1, 1) + mse_loss(fake_d_cLR_1, 0)
        D_loss_cLR_2 = mse_loss(real_d_cLR_2, 1) + mse_loss(fake_d_cLR_2, 0)
        
        D_loss = D_loss_cVAE_1 + D_loss_cLR_1 + D_loss_cVAE_2 + D_loss_cLR_2
        
        # Update
        optim_D_cVAE.zero_grad()
        optim_D_cLR.zero_grad()
        optim_G.zero_grad()
        optim_E.zero_grad()
        D_loss.backward()
        optim_D_cVAE.step()
        optim_D_cLR.step()
        
        ''' ----------------------------- 2. Train G & E ----------------------------- '''
        ############# Step 1. GAN loss to fool discriminator (cVAE_GAN and cLR_GAN) #############
        
        # Encoded latent vector
        mu, log_variance = E(cVAE_data['ground_truth'])
        std = torch.exp(log_variance / 2)
        random_z = util.var(torch.randn(1, params.z_dim))
        encoded_z = (random_z * std) + mu
        
        # Generate fake image and get adversarial loss
        fake_img_cVAE = G(cVAE_data['img'], encoded_z)
        fake_d_cVAE_1, fake_d_cVAE_2 = D_cVAE(fake_img_cVAE)
        
        GAN_loss_cVAE_1 = mse_loss(fake_d_cVAE_1, 1)
        GAN_loss_cVAE_2 = mse_loss(fake_d_cVAE_2, 1)
        
        # Random latent vector
        random_z = util.var(torch.randn(1, params.z_dim))
        
        # Generate fake image and get adversarial loss
        fake_img_cLR = G(cLR_data['img'], random_z)
        fake_d_cLR_1, fake_d_cLR_2 = D_cLR(fake_img_cLR)
        
        GAN_loss_cLR_1 = mse_loss(fake_d_cLR_1, 1)
        GAN_loss_cLR_2 = mse_loss(fake_d_cLR_2, 1)
        
        G_GAN_loss = GAN_loss_cVAE_1 + GAN_loss_cVAE_2 + GAN_loss_cLR_1 + GAN_loss_cLR_2
        
        ############# Step 2. KL-divergence with N(0, 1) (cVAE-GAN) #############
                
        KL_div = params.lambda_kl * torch.sum(0.5 * (mu ** 2 + torch.exp(log_variance) - log_variance - 1))
        
        ############# Step 3. Reconstruction of ground truth image (|G(A, z) - B|) (cVAE-GAN) #############
        img_recon_loss = params.lambda_img * L1_loss(fake_img_cVAE, cVAE_data['ground_truth'])
        
        EG_loss = G_GAN_loss + KL_div + img_recon_loss
        optim_D_cVAE.zero_grad()
        optim_D_cLR.zero_grad()
        optim_G.zero_grad()
        optim_E.zero_grad()
        EG_loss.backward(retain_graph=True)
        optim_E.step()
        optim_G.step()
        
        ''' ----------------------------- 3. Train ONLY G ----------------------------- '''
        ############ Step 1. Reconstrution of random latent code (|E(G(A, z)) - z|) (cLR-GAN) ############
                
        # This step should update ONLY G.
        mu_, log_variance_ = E(fake_img_cLR)
        z_recon_loss = L1_loss(mu_, random_z)
        
        G_alone_loss = params.lambda_z * z_recon_loss
        
        optim_D_cVAE.zero_grad()
        optim_D_cLR.zero_grad()
        optim_G.zero_grad()
        optim_E.zero_grad()
        G_alone_loss.backward()
        optim_G.step()
                        
        # Print error and save intermediate result image and weight
        D_A_losses.append(D_loss.data[0])
        G_GAN_losses.append(G_GAN_loss.data[0])
        KL_div_losses.append(KL_div.data[0])
        img_recon_losses.append(img_recon_loss.data[0])
        G_alone_losses.append(G_alone_loss.data[0])
        
        if iters % params.save_every == 0:
            print('[Epoch : %d / Iters : %d/%d] => D_loss : %f / G_GAN_loss : %f / KL_div : %f / img_recon_loss : %f / z_recon_loss : %f'\
                  %(epoch, iters, len(train_data_loader),D_loss.data[0], G_GAN_loss.data[0], KL_div.data[0], img_recon_loss.data[0], G_alone_loss.data[0]))
            
            # Save intermediate result image
            if os.path.exists(params.result_dir) is False:
                os.makedirs(params.result_dir)
                
            result_img = util.make_img(test_data_loader, G, fixed_z, img_num=params.test_img_num, img_size=params.img_size)
                
            img_name = '{epoch}_{iters}.png'.format(epoch=epoch, iters=iters)
            img_path = os.path.join(params.result_dir, img_name)
            
            torchvision.utils.save_image(result_img, img_path, nrow=params.test_img_num+1)
                
            # Save intermediate weight
            if os.path.exists(params.weight_dir) is False:
                os.makedirs(params.weight_dir)
                
            if epoch is None:
                d_cVAE_name = 'D_cVAE.pkl'
                d_cLR_name = 'D_cLR.pkl'
                g_name = 'G.pkl'
                e_name = 'E.pkl'
            else:
                d_cVAE_name = '{epochs}-{name}'.format(epochs=str(epoch), name='D_cVAE.pkl')
                d_cLR_name = '{epochs}-{name}'.format(epochs=str(epoch), name='D_cLR.pkl')
                g_name = '{epochs}-{name}'.format(epochs=str(epoch), name='G.pkl')
                e_name = '{epochs}-{name}'.format(epochs=str(epoch), name='E.pkl')
                
            torch.save(D_cVAE.state_dict(), os.path.join(params.weight_dir, d_cVAE_name))
            torch.save(D_cVAE.state_dict(), os.path.join(params.weight_dir, d_cLR_name))
            torch.save(G.state_dict(), os.path.join(params.weight_dir, g_name))
            torch.save(E.state_dict(), os.path.join(params.weight_dir, e_name))
    
    D_A_avg_loss = torch.mean(torch.FloatTensor(D_A_losses))
    G_GAN_avg_loss = torch.mean(torch.FloatTensor(G_GAN_losses))
    KL_div_avg_loss = torch.mean(torch.FloatTensor(KL_div_losses))
    img_recon_avg_loss = torch.mean(torch.FloatTensor(img_recon_losses))
    G_alone_avg_loss = torch.mean(torch.FloatTensor(G_alone_losses))
    
    D_A_avg_losses.append(D_A_avg_loss)
    G_GAN_avg_losses.append(G_GAN_avg_loss)
    KL_div_avg_losses.append(KL_div_avg_loss)
    img_recon_avg_losses.append(img_recon_avg_loss)
    G_alone_avg_losses.append(G_alone_avg_loss)
    
    # Save weight at the end of every epoch
    if epoch is None:
        d_cVAE_name = 'D_cVAE.pkl'
        d_cLR_name = 'D_cLR.pkl'
        g_name = 'G.pkl'
        e_name = 'E.pkl'
    else:
        d_cVAE_name = '{epochs}-{name}'.format(epochs=str(epoch), name='D_cVAE.pkl')
        d_cLR_name = '{epochs}-{name}'.format(epochs=str(epoch), name='D_cLR.pkl')
        g_name = '{epochs}-{name}'.format(epochs=str(epoch), name='G.pkl')
        e_name = '{epochs}-{name}'.format(epochs=str(epoch), name='E.pkl')
                
    torch.save(D_cVAE.state_dict(), os.path.join(params.weight_dir, d_cVAE_name))
    torch.save(D_cVAE.state_dict(), os.path.join(params.weight_dir, d_cLR_name))
    torch.save(G.state_dict(), os.path.join(params.weight_dir, g_name))
    torch.save(E.state_dict(), os.path.join(params.weight_dir, e_name))
# -



StarcoderdataJupyterScriptsDedupFiltered