import dataset

import os

import sys

import argparse

import time

import pathlib

import torch

import torch.nn as nn

import torchvision.datasets as datasets

import torchvision.transforms as transforms

import torchvision

from torch.autograd import Variable

import numpy as np

from tqdm import tqdm

from dataset import ImageFilelist

from utils_small import generate_batch_images

import model

from pytorch_ssim import SSIM

from pytorch_msssim import MSSSIM

def weights_init(m):

m.bias.data.fill_(0)

def parse_args():

return args

if __name__ == "__main__":

torch.backends.cudnn.enabled=False

args = parse_args()

#Params

epochs = args.epoch

sample_output = args.sample_output

sample_nums = args.sample_nums

batch_size = args.batch_size

gpu_n = args.gpu

sample_interval = args.sample_interval

dataset_path = args.dataset

torch.cuda.set_device(gpu_n)

learning_rate = 1e-3

dataset = ImageFilelist("./data", dataset_path, transforms.Compose([

#transforms.Grayscale(),

transforms.ToTensor()

]))

data_loader = torch.utils.data.DataLoader(dataset, batch_size, shuffle=True)

D = model.Discriminator(6)

G = model.Generator(5)

D.apply(weights_init)

G.apply(weights_init)

D.cuda()

G.cuda()

print(D)

print(G)

D_criterion = torch.nn.BCEWithLogitsLoss().cuda()

D_optimizer = torch.optim.SGD(D.parameters(), lr=1e-3)

G_criterion = torch.nn.BCEWithLogitsLoss().cuda()

G_l1 = torch.nn.L1Loss().cuda()

G_msssim = MSSSIM().cuda()

G_ssim = SSIM().cuda()

G_optimizer = torch.optim.Adam(G.parameters(), lr=1e-3)

pathlib.Path(sample_output).mkdir(parents=True, exist_ok=True)

pathlib.Path(os.path.join(sample_output, "images")).mkdir(parents=True, exist_ok=True)

d_loss = 0

g_loss = 0

d_to_g_threshold = 0.5

g_to_d_threshold = 0.3

train_d = True

train_g = True

conditional_training = False

_si = 1

for epoch in range(epochs):

for i, (image, label) in enumerate(data_loader):

if conditional_training:

if d_loss - g_loss > d_to_g_threshold:

train_d = True

train_g = False

elif g_loss - d_loss > g_to_d_threshold:

train_g = True

train_d = False

else:

train_d = True

train_g = True

D_optimizer.zero_grad()

G_optimizer.zero_grad()

if train_d:

real_output = D(Variable(image.cuda()))

one_hot_label = torch.from_numpy((np.arange(6) == (label.numpy()[:, None])).astype(float)).float().cuda()

real_loss = D_criterion(real_output, Variable(one_hot_label))

real_loss.backward()

D_optimizer.step()

fake_label = np.array([5 for i in range(len(image))])

G_fake_label = np.random.randint(0, 5, len(image))

fake_label_input = (np.arange(5) == (G_fake_label[:,None])).astype(float)

fake_label_output = (np.arange(6) == (fake_label[:,None])).astype(float)

noise = torch.FloatTensor(len(image), 48, 5, 5).normal_().cuda()

fake_images = G(Variable(torch.from_numpy(fake_label_input).float().cuda()), Variable(noise))

fake_output = D(fake_images.detach())

fake_loss = D_criterion(fake_output, Variable(torch.from_numpy(fake_label_output).float().cuda()))

fake_loss.backward()

D_optimizer.step()

total_loss = real_loss + fake_loss

d_loss = total_loss.data[0]

if train_g:

# Spatial AE

input_label = (np.arange(5) == (label.numpy()[:,None])).astype(float)

input_label = torch.from_numpy(input_label).cuda()

fake_noise = torch.FloatTensor(len(image), 48, 5, 5).normal_().cuda()

ae_images = G(Variable(input_label.float()), Variable(fake_noise))

#l1_loss = G_l1(ae_images, Variable(image.cuda(0)))

ae_loss = G_msssim(ae_images, Variable(image.cuda()))

if np.isnan(torch.mean(ae_loss.data)):

ae_loss = G_ssim(ae_images, Variable(image.cuda()))

#ae_loss = l1_loss - ae_loss

ae_loss.backward()

G_optimizer.step()

fake_label = np.random.randint(0, 5, len(image))

fake_label_input = (np.arange(5) == (fake_label[:,None])).astype(float)

fake_noise = torch.FloatTensor(len(image), 48, 5, 5).normal_().cuda()

generated_images = G(Variable(torch.from_numpy(fake_label_input).float().cuda()), Variable(fake_noise))

generated_output = D(generated_images.detach())

target_label = (np.arange(6) == (fake_label[:,None])).astype(float)

generator_loss = G_criterion(generated_output, Variable(torch.from_numpy(target_label).float().cuda()))

g_loss = generator_loss.data[0]

generator_loss.backward()

G_optimizer.step()

print("Epoch [%d/%d], Iter [%d/%d] D Loss:%.8f, G Loss: %.8f, AE Loss: %.8f" % (epoch + 1, epochs,

i, len(dataset) // batch_size, torch.mean(total_loss.data), torch.mean(generator_loss.data), -torch.mean(ae_loss.data)), end="\r")

print("Epoch [%d/%d], Iter [%d/%d] D Loss:%.8f, G Loss: %.8f, AE Loss: %.8f" % (epoch + 1, epochs,

i, len(dataset) // batch_size, torch.mean(total_loss.data), torch.mean(generator_loss.data), -torch.mean(ae_loss.data)), end="\n")

if _si == sample_interval:

generate_batch_images(G, 5, figure_path=os.path.join(sample_output, "images"), prefix="epoch-%d" % (epoch+1))

model_snapshot = os.path.join(sample_output, "model_snapshot", "epoch-%d" % (epoch + 1))

pathlib.Path(model_snapshot).mkdir(parents=True, exist_ok=True)

torch.save(D.state_dict(), os.path.join(model_snapshot, "D.pkl"))

torch.save(G.state_dict(), os.path.join(model_snapshot, "G.pkl"))

_si = 1

else:

_si += 1