def generate(path_file):
global image_shape
gener = generator(image_shape)
gener.load_weights('with_only_gener_savepoint_5.h5') # load weight
file = Image.open(path_file)
file.load()
array = np.array(file).shape
if (array[2] == 4):
rgb_image = Image.new('RGB', file.size, (255, 255, 255))
rgb_image.paste(file, mask=file.split()[3])
file = rgb_image
print("file size", array)
file = file.resize((128, 128))
file = np.array(file)
print(file.shape)
file = image_procress.normalize(file)
img = np.array([file])
generated_img = gener(img)
generated_img = np.array(generated_img)
term = image_procress.denormalize(generated_img[0])
im = Image.fromarray(term)
return im
def test(random_dim, width, height, channels):
batch_size = 32**2
with tf.variable_scope('input'):
random_input = tf.placeholder(tf.float32,
shape=[None, random_dim],
name='random_input')
fake_image = net.generator(random_input,
channels,
random_dim,
is_train=True)
sess = tf.Session()
saver = tf.train.Saver()
writer = tf.summary.FileWriter('logs/newPokemon/', sess.graph)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
save_path = saver.save(sess, '/tmp/model.ckpt')
ckpt = tf.train.latest_checkpoint('./model/newPokemon')
saver.restore(sess, ckpt)
if not os.path.exists('./test'):
os.makedirs('./test')
sample_noise = np.random.uniform(-1.0, 1.0,
size=[batch_size,
random_dim]).astype(np.float32)
imgtest = sess.run(fake_image, feed_dict={random_input: sample_noise})
final_img = np.vstack(
[np.hstack([imgtest[i * j] for j in range(32)]) for i in range(32)])
save(final_img, './test/mozaic.jpg')
def save_generator_output(self, sess, e, fixed_z, fixed_y):
feed_dict = {self.latent_z: fixed_z, self.inputs_y: fixed_y}
fake_out = sess.run(network.generator(self.latent_z, y=self.inputs_y,
embed_y=True, is_training=False, use_bn=True), feed_dict=feed_dict)
image_fn = os.path.join(self.assets_dir,
'{:s}-{:s}-e{:03d}.png'.format(self.dataset_type, self.gan_loss_type, e + 1))
utils.validation(fake_out, self.val_block_size, image_fn)
return
def __init__(self):
self.generator = lambda x: generator(x, feat_out)
self.discriminator = lambda x: discriminator(x)
self.global_step = tf.train.get_or_create_global_step()
self.train_config()
self.saver = tf.train.Saver(var_list=tf.trainable_variables())
def __init__(self, name, dataset_type, mnist_loader, epochs):
# prepare directories
self.assets_dir = './assets/{:s}'.format(name)
self.ckpt_dir = './checkpoints/{:s}'.format(name)
if not os.path.isdir(self.assets_dir):
os.makedirs(self.assets_dir)
if not os.path.isdir(self.ckpt_dir):
os.makedirs(self.ckpt_dir)
#
self.dataset_type = dataset_type
self.mnist_loader = mnist_loader
# tunable parameters
self.y_dim = 10
self.z_dim = 100
self.learning_rate = 0.0002
self.epochs = epochs
self.batch_size = 128
self.print_every = 30
self.save_every = 1
self.val_block_size = 10
# start building graphs
tf.reset_default_graph()
# create placeholders
self.inputs_x = tf.placeholder(tf.float32, [None, 28, 28, 1],
name='inputs_x')
self.inputs_y = tf.placeholder(tf.float32, [None, self.y_dim],
name='inputs_y')
self.inputs_z = tf.placeholder(tf.float32, [None, self.z_dim],
name='inputs_z')
# create generator & discriminator
self.g_out = network.generator(self.inputs_z,
y=self.inputs_y,
reuse=False,
is_training=True)
self.d_real_logits, _ = network.discriminator(self.inputs_x,
y=self.inputs_y,
reuse=False,
is_training=True)
self.d_fake_logits, _ = network.discriminator(self.g_out,
y=self.inputs_y,
reuse=True,
is_training=True)
# compute model loss
self.d_loss, self.g_loss = self.model_loss(self.d_real_logits,
self.d_fake_logits)
# model optimizer
self.d_opt, self.g_opt = self.model_opt(self.d_loss, self.g_loss)
return
def Main():
real_img = tf.placeholder("float", [BATCH_SIZE, IMG_SIZE, IMG_SIZE, 3])
z = tf.placeholder("float", [BATCH_SIZE, h])
G = generator("generator")
D = discriminator("discriminator")
k_t = tf.get_variable("k", initializer=[0.])
fake_img = G(z, IMG_SIZE, n)
real_logits = D(real_img, IMG_SIZE, n, h)
fake_logits = D(fake_img, IMG_SIZE, n, h)
real_loss = l1_loss(real_img, real_logits)
fake_loss = l1_loss(fake_img, fake_logits)
D_loss = real_loss - k_t * fake_loss
G_loss = fake_loss
M_global = real_loss + tf.abs(GAMMA * real_loss - fake_loss)
global_step = tf.Variable(0, trainable=False)
learning_rate = tf.train.inverse_time_decay(LEARNING_RATE, global_step,
5000, 0.5)
Opt_D = tf.train.AdamOptimizer(learning_rate).minimize(
D_loss, var_list=D.var(), global_step=global_step)
Opt_G = tf.train.AdamOptimizer(learning_rate).minimize(G_loss,
var_list=G.var())
with tf.control_dependencies([Opt_D, Opt_G]):
clip = tf.clip_by_value(k_t + LAMBDA * (GAMMA * real_loss - fake_loss),
0, 1)
update_k = tf.assign(k_t, clip)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
facedata = sio.loadmat("../TrainingSet/facedata.mat")["data"]
saver = tf.train.Saver()
# saver.restore(sess, "./save_para/.\\model.ckpt")
for epoch in range(200):
for i in range(facedata.shape[0] // BATCH_SIZE - 1):
batch = facedata[i * BATCH_SIZE:i * BATCH_SIZE +
BATCH_SIZE, :, :, :] / 127.5 - 1.0
z0 = np.random.uniform(0, 1, [BATCH_SIZE, h])
sess.run(update_k, feed_dict={real_img: batch, z: z0})
if i % 100 == 0:
[dloss, gloss, Mglobal, fakeimg, step,
lr] = sess.run([
D_loss, G_loss, M_global, fake_img, global_step,
learning_rate
],
feed_dict={
real_img: batch,
z: z0
})
print(
"step: %d, d_loss: %f, g_loss: %f, M_global: %f, Learning_rate: %f"
% (step, dloss, gloss, Mglobal, lr))
Image.fromarray(np.uint8(
127.5 * (fakeimg[0, :, :, :] + 1))).save("./Results/" +
str(step) +
".jpg")
saver.save(sess, "./save_para/model.ckpt")
def up_scale(downsampled_img):
downsampled_img = downsampled_img[np.newaxis, :, :, :]
downsampled = tf.placeholder(tf.float32, [None, None, None, 3])
train_phase = tf.placeholder(tf.bool)
G = generator("generator")
SR = G(downsampled, train_phase)
sess = tf.Session()
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(sess, "./save_para/.\\model.ckpt")
SR_img = sess.run(SR, feed_dict={downsampled: downsampled_img/127.5 - 1, train_phase: False})
Image.fromarray(np.uint8((SR_img[0, :, :, :] + 1)*127.5)).show()
Image.fromarray(np.uint8((downsampled_img[0, :, :, :]))).show()
sess.close()
def test_generator(self):
n = 2
h = 128
w = h
c = 4
class_num = 3
input_tensor = tf.random_uniform((n, h, w, c))
target_tensor = tf.random_uniform((n, class_num))
output_tensor = network.generator(input_tensor, target_tensor)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
output = sess.run(output_tensor)
self.assertTupleEqual((n, h, w, c), output.shape)
def test_generator(self):
n = 2
h = 128
w = h
c = 4
class_num = 3
input_tensor = tf.random_uniform((n, h, w, c))
target_tensor = tf.random_uniform((n, class_num))
output_tensor = network.generator(input_tensor, target_tensor)
with self.test_session() as sess:
sess.run(tf.global_variables_initializer())
output = sess.run(output_tensor)
self.assertTupleEqual((n, h, w, c), output.shape)
def test_dehaze(args, logger=None):
batch_size = args.batch_size
num_workers = args.workers
phase = args.phase
for k, v in args.__dict__.items():
print(k, ':', v)
model = generator(3, 3)
model = nn.DataParallel(model).cuda()
data_dir = args.data_dir
dataset = DehazeList(
data_dir,
phase,
transforms.Compose([
transforms.ToTensor(),
#normalize,
]),
out_name=True)
test_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=False)
cudnn.benchmark = True
# optionally resume from a checkpoint
start_epoch = 0
if args.resume:
if os.path.isfile(args.resume):
logger.info("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
best_psnr1 = checkpoint['best_psnr1']
model.load_state_dict(checkpoint['state_dict'])
logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
logger.info("=> no checkpoint found at '{}'".format(args.resume))
return
out_dir = '{:03d}_{}'.format(start_epoch, phase)
test(test_loader, model, save_vis=True, output_dir=out_dir, logger=logger)
def train(batch_size=4, lambd=1e-3, init_lr=1e-4, clip_v=0.05, B=16, max_itr=100000, path_trainset="./ImageNet/", path_vgg="./vgg_para/", path_save_model="./save_para/"):
inputs = tf.placeholder(tf.float32, [None, 96, 96, 3])
downsampled = tf.placeholder(tf.float32, [None, 24, 24, 3])
train_phase = tf.placeholder(tf.bool)
learning_rate = tf.placeholder(tf.float32)
G = generator("generator", B)
D = discriminator("discriminator")
SR = G(downsampled, train_phase)
phi = vggnet(tf.concat([inputs, SR], axis=0), path_vgg)
phi = tf.split(phi, num_or_size_splits=2, axis=0)
phi_gt = phi[0]
phi_sr = phi[1]
real_logits = D(inputs, train_phase)
fake_logits = D(SR, train_phase)
D_loss = tf.reduce_mean(fake_logits) - tf.reduce_mean(real_logits)
G_loss = -tf.reduce_mean(fake_logits) * lambd + tf.nn.l2_loss(phi_sr - phi_gt) / batch_size
clip_D = [var.assign(tf.clip_by_value(var, -clip_v, clip_v)) for var in D.var_list()]
D_opt = tf.train.RMSPropOptimizer(learning_rate).minimize(D_loss, var_list=D.var_list())
G_opt = tf.train.RMSPropOptimizer(learning_rate).minimize(G_loss, var_list=G.var_list())
sess = tf.Session()
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
# saver.restore(sess, "./save_para/.\\model.ckpt")
lr0 = init_lr
for itr in range(max_itr):
if itr == max_itr // 2 or itr == max_itr * 3 // 4:
lr0 = lr0 / 10
s0 = time.time()
batch, down_batch = read_crop_data(path_trainset, batch_size, [96, 96, 3], 4)
e0 = time.time()
batch = batch/127.5 - 1
down_batch = down_batch/127.5 - 1
s1 = time.time()
sess.run(D_opt, feed_dict={inputs: batch, downsampled: down_batch, train_phase: True, learning_rate: lr0})
sess.run(clip_D)
sess.run(G_opt, feed_dict={inputs: batch, downsampled: down_batch, train_phase: True, learning_rate: lr0})
e1 = time.time()
if itr % 200 == 0:
[d_loss, g_loss, sr] = sess.run([D_loss, G_loss, SR], feed_dict={downsampled: down_batch, inputs: batch, train_phase: False})
raw = np.uint8((batch[0] + 1) * 127.5)
bicub = misc.imresize(np.uint8((down_batch[0] + 1) * 127.5), [96, 96])
gen = np.uint8((sr[0, :, :, :] + 1) * 127.5)
print("Iteration: %d, D_loss: %f, G_loss: %e, PSNR: %f, SSIM: %f, Read_time: %f, Update_time: %f" % (itr, d_loss, g_loss, psnr(raw, gen), ssim(raw, gen, multichannel=True), e0 - s0, e1 - s1))
Image.fromarray(np.concatenate((raw, bicub, gen), axis=1)).save("./results/" + str(itr) + ".jpg")
if itr % 5000 == 0:
saver.save(sess, path_save_model+"model.ckpt")
def __init__(self):
self.dataloader = initialize_dataloader()
self.generator = network.generator().to(configs.device)
self.discriminator = network.discriminator().to(configs.device)
self.optimizer_generator = torch.optim.Adam(
self.generator.parameters(),
lr=configs.learning_rate,
betas=configs.betas)
self.optimizer_discriminator = torch.optim.Adam(
self.discriminator.parameters(),
lr=configs.learning_rate,
betas=configs.betas)
self.loss = torch.nn.BCELoss()
self.disloss = []
self.genloss = []
def __init__(self, training=True):
N.bn_training = training
self.inpholder = tf.placeholder(tf.float32, [None, 128, 128, 3])
self.feat = N.feat_encoder(self.inpholder)
self.age_features = []
self.recon_features = []
self.recon_imgs = []
self.dis_scores = []
self.dis_feat_scores = []
self.dis_scores_real = []
self.dis_feat_scores_real = []
for i in range(10):
age_feat = N.age_encoder(self.feat, i)
self.age_features.append(age_feat)
recon_img = N.generator(age_feat)
self.recon_imgs.append(recon_img)
recon_feature = N.feat_encoder(recon_img)
self.recon_features.append(recon_feature)
dis_scr = N.discriminator(recon_img)
self.dis_scores.append(dis_scr)
feat_dis_scr = N.discriminator_feature(recon_feature)
self.dis_feat_scores.append(feat_dis_scr)
dis_scr_real = N.discriminator(self.inpholder)
self.dis_scores_real.append(dis_scr_real)
feat_dis_scr_real = N.discriminator_feature(self.feat)
self.dis_feat_scores_real.append(feat_dis_scr_real)
self.age_features = tf.stack(self.age_features, 1)
self.img_feature = N.attention_blk(self.age_features)
def test(downsampled_img, img, B):
downsampled_img = downsampled_img[np.newaxis, :, :, :]
downsampled = tf.placeholder(tf.float32, [None, None, None, 3])
train_phase = tf.placeholder(tf.bool)
G = generator("generator", B)
SR = G(downsampled, train_phase)
sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(sess, "./save_para/.\\model.ckpt")
SR_img = sess.run(SR, feed_dict={downsampled: downsampled_img/127.5 - 1, train_phase: False})
Image.fromarray(np.uint8((SR_img[0, :, :, :] + 1)*127.5)).show()
Image.fromarray(np.uint8((downsampled_img[0, :, :, :]))).show()
h = img.shape[0]
w = img.shape[1]
bic_img = misc.imresize(downsampled_img[0, :, :, :], [h, w])
Image.fromarray(np.uint8((bic_img))).show()
SR_img = misc.imresize(SR_img[0, :, :, :], [h, w])
p = psnr(img, SR_img)
s = ssim(img, SR_img, multichannel=True)
p1 = psnr(img, bic_img)
s1 = ssim(img, bic_img, multichannel=True)
print("SR PSNR: %f, SR SSIM:%f, BIC PSNR: %f, BIC SSIM: %f"%(p, s, p1, s1))
sess.close()
def __init__(self, name, dataset_type, gan_loss_type):
# prepare directories
self.assets_dir = './assets/{:s}'.format(name)
self.ckpt_dir = './ckpts/{:s}'.format(name)
self.ckpt_fn = os.path.join(self.ckpt_dir, '{:s}.ckpt'.format(name))
if not os.path.exists(self.assets_dir):
os.makedirs(self.assets_dir)
if not os.path.exists(self.ckpt_dir):
os.makedirs(self.ckpt_dir)
# setup variables
self.dataset_type = dataset_type
# tunable parameters
self.z_dim = 100
self.learning_rate = 5e-5
self.epochs = 30
self.batch_size = 128
self.print_every = 30
self.save_every = 5
self.val_block_size = 10
# start building graphs
tf.reset_default_graph()
# create placeholders
self.latent_z = tf.placeholder(tf.float32, [None, self.z_dim],
name='latent_z')
self.real_images = tf.placeholder(tf.float32, [None, 28, 28, 1],
name='real_images')
# create generator & discriminator
self.fake_images = network.generator(self.latent_z,
is_training=True,
use_bn=True)
self.d_real_logits, _ = network.discriminator(self.real_images,
is_training=True,
use_bn=True)
self.d_fake_logits, _ = network.discriminator(self.fake_images,
is_training=True,
use_bn=True)
# compute model loss
self.d_loss, self.g_loss = wgan_loss(self.d_real_logits,
self.d_fake_logits)
# prepare optimizers
t_vars = tf.trainable_variables()
d_vars = [
var for var in t_vars if var.name.startswith('discriminator')
]
g_vars = [var for var in t_vars if var.name.startswith('generator')]
# add clipping op
self.d_clip_op = tf.group(
[p.assign(tf.clip_by_value(p, -0.01, 0.01)) for p in d_vars])
# Optimize
optimizer = tf.train.RMSPropOptimizer(self.learning_rate)
with tf.control_dependencies(tf.get_collection(
tf.GraphKeys.UPDATE_OPS)):
self.d_opt = optimizer.minimize(self.d_loss, var_list=d_vars)
self.g_opt = optimizer.minimize(
self.g_loss,
var_list=g_vars,
global_step=tf.train.get_or_create_global_step())
# prepare saver for generator
self.saver = tf.train.Saver(var_list=g_vars)
return
def validate(val_loader, model, criterion, print_freq=10, output_dir='val', \
save_vis=False, epoch=None, eval_score=None, logger=None, auto_save=True, best_score=0.0):
batch_time = AverageMeter()
losses = AverageMeter()
score = AverageMeter()
# switch to evaluate mode
generator = model
generator.eval()
end = time.time()
for i, (input, target, name) in enumerate(val_loader):
input = input.float()
target = target.float()
input = input.cuda()
target = target.cuda(async=True)
input_var = torch.autograd.Variable(input)
target_var = torch.autograd.Variable(target)
# compute output
with torch.no_grad():
output, _, _ = generator(input_var)
loss = criterion(output, target_var)
losses.update(loss.data, input.size(0))
if eval_score is not None:
score.update(eval_score(output, target_var), input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
save_dir = os.path.join(output_dir, 'epoch_{:04d}'.format(epoch + 1))
if save_vis == True:
save_dir = os.path.join(output_dir,
'epoch_{:04d}'.format(epoch + 1))
pred = output
save_output_images(pred, name, save_dir)
if auto_save == True and (score.avg) > best_score:
save_dir = os.path.join(output_dir,
'best/' + 'epoch_{:04d}'.format(epoch + 1))
pred = output
save_output_images(pred, name, save_dir)
logger.info('Best model: {0}'.format(epoch + 1))
if i % print_freq == 0:
logger.info('Test: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Score {top1.val:.3f} ({top1.avg:.3f})'.format(
i,
len(val_loader),
batch_time=batch_time,
loss=losses,
top1=score))
logger.info(' * Score {top1.avg:.3f}'.format(top1=score))
print()
return score.avg
def train_dehaze(args, saveDirName='.', logger=None):
batch_size = args.batch_size
num_workers = args.workers
crop_size = args.crop_size
print(' '.join(sys.argv))
# logging hyper-parameters
for k, v in args.__dict__.items():
logger.info('{0}:\t{1}'.format(k, v))
# Generators
net = generator(3, 3)
net = nn.DataParallel(net).cuda()
model = net
# Criterion for updating weights
criterion = nn.L1Loss()
criterion = criterion.cuda()
# Data loading code
data_dir = args.data_dir
t = []
if args.random_scale > 0:
t.append(transforms.RandomScale(args.random_scale))
t.append(transforms.RandomCrop(crop_size))
if args.random_rotate > 0:
t.append(transforms.RandomRotate(args.random_rotate))
t.extend([
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.RandomIdentityMapping(p=0.4),
transforms.ToTensor(),
])
# DataLoaders for training/validation dataset
train_loader = torch.utils.data.DataLoader(DehazeList(
data_dir, 'train', transforms.Compose(t), out_name=False),
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
pin_memory=True,
drop_last=True)
val_loader = torch.utils.data.DataLoader(DehazeList(
data_dir,
'val',
transforms.Compose([
transforms.ToTensor(),
]),
out_name=True),
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=False,
drop_last=False)
# define loss function (criterion) and optimizer
optimizer = torch.optim.Adam(net.parameters(),
args.lr,
betas=(0.5, 0.999),
weight_decay=args.weight_decay)
cudnn.benchmark = True
best_psnr1 = 0
start_epoch = 0
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
best_psnr1 = checkpoint['best_psnr1']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
lr = args.lr
for epoch in range(start_epoch, args.epochs):
lr = adjust_learning_rate(args, optimizer, epoch, lr)
logger.info('Epoch: [{0}]\tlr {1:.06f}'.format(epoch + 1, lr))
train(train_loader,
model,
criterion,
optimizer,
epoch,
eval_score=psnr,
logger=logger)
psnr1 = 0
if epoch % 5 == 4:
psnr1 = validate(val_loader, model, criterion, eval_score=psnr, save_vis=True, \
epoch=epoch, logger=logger, best_score=best_psnr1)
else:
psnr1 = validate(val_loader, model, criterion, eval_score=psnr, epoch=epoch, \
logger=logger, best_score=best_psnr1)
if epoch == 0:
best_psnr1 = psnr1
is_best = (psnr1 >= best_psnr1)
best_psnr1 = max(psnr1, best_psnr1)
checkpoint_path = saveDirName + '/' + 'checkpoint_latest.pth.tar'
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_psnr1': best_psnr1,
},
is_best,
filename=checkpoint_path)
if (epoch + 1) % 1 == 0:
history_path = saveDirName + '/' + 'checkpoint_{:03d}.pth.tar'.format(
epoch + 1)
shutil.copyfile(checkpoint_path, history_path)
def main(_):
if not os.path.exists(FLAGS.checkpoint_dir):
os.makedirs(FLAGS.checkpoint_dir)
if not os.path.exists(FLAGS.sample_dir):
os.makedirs(FLAGS.sample_dir)
if not os.path.exists(FLAGS.summaries_dir):
os.makedirs(FLAGS.summaries_dir)
with tf.device("/gpu:0"):
#with tf.device("/cpu:0"):
z = tf.placeholder(tf.float32, [FLAGS.batch_size, FLAGS.z_dim], name="g_input_noise")
x = tf.placeholder(tf.float32, [FLAGS.batch_size, FLAGS.output_size, FLAGS.output_size, FLAGS.c_dim], name='d_input_images')
Gz = network.generator(z)
Dx, Dfx = network.discriminator(x)
Dz, Dfz = network.discriminator(Gz, reuse=True)
d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=Dx, labels=tf.ones_like(Dx)))
d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=Dz, labels=tf.zeros_like(Dz)))
d_loss = d_loss_real + d_loss_fake
g_loss_perceptual = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits = Dz, labels = tf.ones_like(Dz)))
g_loss_features = tf.reduce_mean(tf.nn.l2_loss(Dfx-Dfz))/(FLAGS.image_size*FLAGS.image_size)
g_loss = g_loss_perceptual + g_loss_features
tvars = tf.trainable_variables()
d_vars = [var for var in tvars if 'd_' in var.name]
g_vars = [var for var in tvars if 'g_' in var.name]
print(d_vars)
print("---------------")
print(g_vars)
with tf.variable_scope(tf.get_variable_scope(),reuse=False):
print("reuse or not: {}".format(tf.get_variable_scope().reuse))
assert tf.get_variable_scope().reuse == False, "Houston tengo un problem"
d_trainer = tf.train.AdamOptimizer(FLAGS.learning_rate, FLAGS.beta1).minimize(d_loss, var_list=d_vars)
g_trainer = tf.train.AdamOptimizer(FLAGS.learning_rate, FLAGS.beta1).minimize(g_loss, var_list=g_vars)
tf.summary.scalar("generator_loss_percptual", g_loss_perceptual)
tf.summary.scalar("generator_loss_features", g_loss_features)
tf.summary.scalar("generator_loss_total", g_loss)
tf.summary.scalar("discriminator_loss", d_loss)
tf.summary.scalar("discriminator_loss_real", d_loss_real)
tf.summary.scalar("discriminator_loss_fake", d_loss_fake)
images_for_tensorboard = network.generator(z, reuse=True)
tf.summary.image('Generated_images', images_for_tensorboard, 2)
merged = tf.summary.merge_all()
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.30)
gpu_options.allow_growth = True
saver = tf.train.Saver()
with tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, allow_soft_placement=True)) as sess:
print("starting session")
summary_writer = tf.summary.FileWriter(FLAGS.summaries_dir + '/train', sess.graph)
sess.run(tf.global_variables_initializer())
data_files = glob(os.path.join("./data", FLAGS.dataset, "*.jpg"))
model_dir = "%s_%s_%s" % (FLAGS.dataset, 64, FLAGS.output_size)
save_dir = os.path.join(FLAGS.checkpoint_dir, model_dir)
if FLAGS.is_train:
for epoch in range(FLAGS.epoch):
d_total_cost = 0.
g_total_cost = 0.
shuffle(data_files)
num_batches = min(len(data_files), FLAGS.train_size) // FLAGS.batch_size
#num_batches = 2
for batch_i in range(num_batches):
batch_files = data_files[batch_i*FLAGS.batch_size:(batch_i+1)*FLAGS.batch_size]
batch = [utilities.load_image(batch_file, FLAGS.image_size, is_crop=FLAGS.is_crop, resize_w=FLAGS.output_size) for batch_file in batch_files]
batch_x = np.array(batch).astype(np.float32)
batch_z = np.random.normal(-1, 1, size=[FLAGS.batch_size, FLAGS.z_dim]).astype(np.float32)
start_time = time.time()
d_err, _ = sess.run([d_loss, d_trainer], feed_dict={z: batch_z, x: batch_x})
g_err, _ = sess.run([g_loss, g_trainer], feed_dict={z: batch_z, x: batch_x})
d_total_cost += d_err
g_total_cost += g_err
if batch_i % 10 == 0:
summary = sess.run(merged, feed_dict={x: batch_x, z: batch_z})
summary_writer.add_summary(summary, (epoch-1)*(num_batches/30)+(batch_i/30))
print("Epoch: [%2d/%2d] [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f" \
% (epoch, FLAGS.epoch, batch_i, num_batches,
time.time() - start_time, d_err, g_err))
print("Epoch:", '%04d' % (epoch+1), "d_cost=", \
"{:.9f}".format(d_total_cost/num_batches), "g_cost=", "{:.9f}".format(g_total_cost/num_batches))
sys.stdout.flush()
save_path = saver.save(sess, save_dir)
print("Model saved in path: %s" % save_path)
sys.stdout.flush()
sess.close()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', required=True, help='cifar10 | imagenet')
parser.add_argument('--dataroot', default='./data/', required=True, help='path to dataset')
parser.add_argument('--out', help='Directory to output the result')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ndf', type=int, default=64)
parser.add_argument('--batchsize', type=int, default=32, help='input batch size')
parser.add_argument('--imagesize', default=64, help='the height / width of the input image to network')
parser.add_argument('--lr', type=float, default=0.0002, help='learning rate')
parser.add_argument('--nz', default=64, help='Number of hidden units(z)')
parser.add_argument('--epochs', default=100, help='number of epochs to train for')
parser.add_argument('--ngpu', default=1, help='number of GPUs to use')
parser.add_argument('--cuda', help='enables cuda')
opt=parser.parse_args()
print(opt)
data_loader = CIFAR(batch_size= opt.batchsize)
nc = 3
'''
if opt.dataset == 'cifar10':
dataset = dset.CIFAR10(root=opt.dataroot, download=True,
transform=transforms.Compose([
transforms.Resize(opt.imagesize),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]))
nc = 3
else:
dataset = dset.ImageFolder(root=opt.dataroot,
transform=transforms.Compose([
transforms.Resize(opt.imagesize),
transforms.CenterCrop(opt.imagesize),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]))
'''
'''
GPUの使用について（例）
x = torch.randn(10)
y = torch.randn(10)
x = x.to('cuda')
y = y.to('cuda:0') # cudaの後に：数字で対応したGPUを使用
z = x * y
z = z.to('cpu') # cpuへ
'''
G_net = generator(ngf=opt.ngf, ngpu=1, nc=3, nz=opt.nz).to(device)
#D_net = discriminator(ndf=opt.ndf, nc=3, ngpu=0).to(device)
D_net = discriminator(ndf=opt.ndf, ngpu=1, nc=nc).to(device)
##################
### ここから上までで、モデルの構築
##################
# Initialize BCELoss function
criterion = nn.BCELoss().to(device)
# assert dataset
# dataset = dset.ImageFolder(root = opt.dataroot)
# dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchsize, shuffle=True)
nc = 3
# setup optimizer
optimizerG = torch.optim.Adam(G_net.parameters(), lr=opt.lr)
optimizerD = torch.optim.Adam(D_net.parameters(), lr=opt.lr, betas=(0.5, 0.999))
# Generate batch of latent vectors
input_noise = torch.randn(opt.batchsize, opt.nz, 1, 1, device=device) # batch * channels * height * width
# Establish convention for real and fake labels during training
real_label = 1
fake_label = 0
# Lists to keep track of progress
G_loss_list = []
D_loss_list = []
print("Starting training loop...")
'''
＜DCGANのアルゴリズム＞
1.ミニバッチサイズm個のノイズz1, z2, ..., zmをPg(z)から取り出す(生成する)
2.ミニバッチサイズm個のサンプルx1, x2, ..., xmをデータ生成分布Pdata(x)から取り出す
3.1/m sigma((log(D(x))) + (log(1 - D(G(z)))))式の、theta(d)における確率的勾配を上るようにDを更新
4.上記までをk回くりかえす
5.ミニバッチサイズm個のノイズz1, z2, ..., zmをPg(z)から取り出す
6.1/m sigma(log(1 - D(G(z)))))式の、theta(g)における確率的勾配を下るようにGを更新
7.ここまで全てを、訓練回数分だけ繰り返す
(8.)Dを十分な回数(k回)更新した上で、Gを1回更新することで、常に鑑別機が新しいGの状態に適用できるように学習を進める
'''
# For each epoch
for epoch in range(opt.epochs):
# For each batch in the dataloader
for i_, data in enumerate(data_loader):
#######################
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z))
#######################
## Train with all-real batch
# configure training data
real_img = Variable(data[0]).to(device)
# 勾配の初期化
optimizerD.zero_grad()
# Forward pass through D
real_label = Variable(torch.ones(opt.batchsize)).to(device)
fake_label = Variable(torch.zeros(opt.batchsize)).to(device)
real_out = D_net(real_img)
# Calculate loss on all-real batch
d_loss_real = criterion(real_out, real_label)
## Train with all-fake batch
z = Variable(torch.randn(opt.batchsize, opt.nz, 1, 1)).to(device)
fake_img = G_net(z)
fake_out = D_net(fake_img)
d_loss_fake = criterion(fake_out, fake_label)
# Add the gradients from the all-real and all-fake batches
d_loss = d_loss_real + d_loss_fake
# Calculate gradients for D in backward pass
d_loss.backward()
# Update D
optimizerD.step()
#######################
# (2) Update G network: maximize log(D(G(z))
#######################
G_net.zero_grad()
fake_img = G_net(z)
fake_out = D_net(fake_img)
g_loss = criterion(fake_out, real_label)
# Calculate gradients for G in backward pass
g_loss.backward()
optimizerG.step()
i = 0
# Output training stats
if i % 1000 == 0:
print(fake_img.size())
# Save Losses for plotting later
G_loss_list.append(g_loss.item())
D_loss_list.append(d_loss.item())
def train(dir, random_dim, width, height, channels, batch_size, epoch):
with tf.variable_scope('input'):
real_image = tf.placeholder(tf.float32, shape = [None, height, width, channels], name = 'real_image')
random_input = tf.placeholder(tf.float32, shape = [None, random_dim], name = 'random_input')
fake_image = net.generator(random_input, channels, random_dim, is_train = True)
real_result, _ = net.discriminator(real_image, is_train = True)
fake_result, _ = net.discriminator(fake_image, is_train = True, reuse = True)
fake_result_mean = tf.reduce_mean(fake_result)
d_loss = tf.reduce_mean(real_result) - fake_result_mean
g_loss = -fake_result_mean
t_vars = tf.trainable_variables()
d_vars = [var for var in t_vars if 'discriminator' in var.name]
g_vars = [var for var in t_vars if 'generator' in var.name]
learning_rate = 1e-3
trainer_d = tf.train.AdamOptimizer(learning_rate).minimize(-d_loss, var_list = d_vars)
trainer_g = tf.train.AdamOptimizer(learning_rate).minimize(g_loss, var_list = g_vars)
d_clip = [v.assign(tf.clip_by_value(v, -0.01, 0.01)) for v in d_vars]
images_batch, samples_num = net.get_images_batch(dir, width, height, channels, batch_size)
batch_num = int(samples_num / batch_size)
total_batch = 0
sess = tf.Session()
saver = tf.train.Saver()
writer = tf.summary.FileWriter('logs/newPokemon/', sess.graph)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
save_path = saver.save(sess, '/tmp/model.ckpt')
ckpt = tf.train.latest_checkpoint('./model/newPokemon')
saver.restore(sess, ckpt)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess = sess, coord = coord)
tf.summary.scalar('loss_discriminator', d_loss)
tf.summary.scalar('loss_generator', g_loss)
summary_op = tf.summary.merge_all()
print('total training sample num: %d' % samples_num)
print('batch size: %d, batch num per epoch: %d, epoch num: %d' % (batch_size, batch_num, epoch))
print('start training...')
for i in tqdm(range(epoch)):
for j in range(batch_num):
d_iters = 5
g_iters = 1
train_noise = np.random.uniform(-1.0, 1.0, size = [batch_size, random_dim]).astype(np.float32)
for k in range(d_iters):
train_image = sess.run(images_batch)
sess.run(d_clip)
_, dLoss = sess.run([trainer_d, d_loss], feed_dict = {random_input: train_noise, real_image: train_image})
for k in range(g_iters):
_, gLoss = sess.run([trainer_g, g_loss], feed_dict = {random_input: train_noise})
if i == 0:
if not os.path.exists('./newPokemon'):
os.makedirs('./newPokemon')
for index in range(train_image.shape[0]):
image = train_image[index]
save(image, './newPokemon/batch' + str(i) + '_image' + str(index) + '.jpg')
if i % 100 == 0:
if not os.path.exists('./model/newPokemon'):
os.makedirs('./model/newPokemon')
saver.save(sess, './model/newPokemon/' + str(i))
if i % 50 == 0:
if not os.path.exists('./newPokemon'):
os.makedirs('./newPokemon')
sample_noise = np.random.uniform(-1.0, 1.0, size = [10, random_dim]).astype(np.float32) #[batch_size, random_dim]).astype(np.float32)
imgtest = sess.run(fake_image, feed_dict = {random_input: sample_noise})
for index in range(imgtest.shape[0]):
image = imgtest[index]
save(image, './newPokemon/epoch' + str(i) + '_image' + str(index) + '.jpg')
summary_str = sess.run(summary_op, feed_dict = {random_input: train_noise, real_image: train_image})
writer.add_summary(summary_str, i)
print('train:[%d],d_loss:%f,g_loss:%f' % (i, dLoss, gLoss))
coord.request_stop()
coord.join(threads)
num_workers=4)
### training phase
# parameters setting
# FOR HA
B_real_ = torch.randn(num_train, opt.bit)
B_fake_ = torch.randn(num_train, opt.bit)
H_I_ = torch.zeros(num_train, opt.bit)
H_S_ = torch.zeros(num_train, opt.bit)
B_I = torch.sign(torch.sign(B_real_))
B_S = torch.sign(torch.sign(B_fake_))
# network
G_A = network.generator(opt.input_ngc, opt.output_ngc, opt.ngf, opt.nb)
G_B = network.generator(opt.input_ngc, opt.output_ngc, opt.ngf, opt.nb)
D_A = network.discriminator(opt.input_ndc, opt.output_ngc, opt.ndf)
net = models.resnet18(pretrained=True)
net.fc = nn.Linear(2048, opt.bit)
# net_dict = net.state_dict()
# pretrain_dict = torch.load('./pre_resnet18.pkl')
# pretrain_dict = {k[7:] : v for k, v in pretrain_dict.items() if k[7:] in net_dict}
# net_dict.update(pretrain_dict)
# net.load_state_dict(net_dict)
# net.fc= nn.Linear(2048, opt.bit)
H_A = net
test_loader = torch.utils.data.DataLoader(MNIST(root='./data/mnist',
train=False),
batch_size=128,
shuffle=False,
num_workers=4,
pin_memory=True)
vggish_model = Vggish().to(device)
load_weight(vggish_model, './save_models/vggish_mnist_best.pth')
class_model = ClassifyNet().to(device)
class_model.load_state_dict(
torch.load('./save_models/classify_mnist_best.pth')['shared_layers'])
G_model = generator(out_size=3).to(device)
G_model.load_state_dict(torch.load('./save_models/emnist_G_best.pth.tar'))
D_model = discriminator(in_size=3, ndf=128).to(device)
D_model.load_state_dict(torch.load('./save_models/emnist_D_best.pth.tar'))
optimizer = optim.Adam(vggish_model.parameters(), lr=1e-3)
prec = test(vggish_model, class_model, G_model, test_loader, 0)
best_prec = prec
for epoch in range(1, 5):
if epoch > 3:
optimizer = optim.Adam(list(vggish_model.parameters()) +
list(G_model.parameters()),
lr=1e-4)
train(vggish_model, class_model, G_model, D_model, train_loader,
def __init__(self, name, dataset_type, gan_loss_type):
# prepare directories
self.assets_dir = './assets/{:s}'.format(name)
self.ckpt_dir = './ckpts/{:s}'.format(name)
self.ckpt_fn = os.path.join(
self.ckpt_dir, '{:s}-{:s}.ckpt'.format(name, gan_loss_type))
if not os.path.exists(self.assets_dir):
os.makedirs(self.assets_dir)
if not os.path.exists(self.ckpt_dir):
os.makedirs(self.ckpt_dir)
# setup variables
self.dataset_type = dataset_type
self.gan_loss_type = gan_loss_type
# tunable parameters
self.y_dim = 10
self.z_dim = 128
self.learning_rate = 1e-4
self.epochs = 30
self.batch_size = 128
self.print_every = 30
self.save_every = 5
self.val_block_size = 10
# start building graphs
tf.reset_default_graph()
# create placeholders
self.latent_z = tf.placeholder(tf.float32, [None, self.z_dim],
name='latent_z')
self.inputs_y = tf.placeholder(tf.float32, [None, self.y_dim],
name='inputs_y')
self.real_images = tf.placeholder(tf.float32, [None, 28, 28, 1],
name='real_images')
# create generator & discriminator
self.fake_images = network.generator(self.latent_z,
y=self.inputs_y,
is_training=True,
use_bn=True)
self.d_real_logits, self.a_real_input = network.discriminator(
self.real_images, y=self.inputs_y, is_training=True, use_bn=True)
self.d_fake_logits, self.a_fake_input = network.discriminator(
self.fake_images, y=self.inputs_y, is_training=True, use_bn=True)
self.a_real_logits = network.classifier(self.a_real_input,
self.y_dim,
is_training=True,
use_bn=True)
self.a_fake_logits = network.classifier(self.a_fake_input,
self.y_dim,
is_training=True,
use_bn=True)
# compute model loss
if gan_loss_type == 'v1':
self.d_loss, self.g_loss = gan_loss_v1(self.d_real_logits,
self.d_fake_logits)
elif gan_loss_type == 'v2':
self.d_loss, self.g_loss = gan_loss_v2(self.d_real_logits,
self.d_fake_logits)
else:
raise ValueError('gan_loss_type must be either "v1" or "v2"!!')
self.a_loss = auxilary_classifier_loss(self.a_real_logits,
self.a_fake_logits,
self.inputs_y)
# prepare optimizers
t_vars = tf.trainable_variables()
d_vars = [
var for var in t_vars if var.name.startswith('discriminator')
]
g_vars = [var for var in t_vars if var.name.startswith('generator')]
# Optimize
optimizer = tf.train.AdamOptimizer(self.learning_rate, beta1=0.5)
with tf.control_dependencies(tf.get_collection(
tf.GraphKeys.UPDATE_OPS)):
self.d_opt = optimizer.minimize(self.d_loss, var_list=d_vars)
self.g_opt = optimizer.minimize(self.g_loss, var_list=g_vars)
self.a_opt = optimizer.minimize(
self.a_loss,
var_list=t_vars,
global_step=tf.train.get_or_create_global_step())
# prepare saver for generator
self.saver = tf.train.Saver(var_list=g_vars)
return
def train(self):
val_size = self.val_block_size * self.val_block_size
steps = 0
losses = []
new_epochs = self.d_train_freq * self.epochs
start_time = time.time()
with tf.Session() as sess:
# reset tensorflow variables
sess.run(tf.global_variables_initializer())
# start training
for e in range(new_epochs):
for ii in range(self.mnist_loader.train.num_examples //
self.batch_size):
# no need labels
batch_x, _ = self.mnist_loader.train.next_batch(
self.batch_size)
# rescale images to -1 ~ 1
batch_x = np.reshape(batch_x, (-1, 28, 28, 1))
batch_x = batch_x * 2.0 - 1.0
# Sample random noise for G
batch_z = np.random.uniform(-1,
1,
size=(self.batch_size,
self.z_dim))
fd = {self.inputs_x: batch_x, self.inputs_z: batch_z}
# Run optimizers (train D more than G)
_ = sess.run(self.d_weight_clip)
_ = sess.run(self.d_opt, feed_dict=fd)
if ii % self.d_train_freq == 0:
_ = sess.run(self.g_opt, feed_dict=fd)
# print losses
if steps % self.print_every == 0:
# At the end of each epoch, get the losses and print them out
train_loss_d = self.d_loss.eval({
self.inputs_x: batch_x,
self.inputs_z: batch_z
})
train_loss_g = self.g_loss.eval(
{self.inputs_z: batch_z})
print(
"Epoch {}/{}...".format(e + 1, self.epochs),
"Discriminator Loss: {:.4f}...".format(
train_loss_d),
"Generator Loss: {:.4f}".format(train_loss_g))
losses.append((train_loss_d, train_loss_g))
steps += 1
# save generation results at every epochs
if e % (self.d_train_freq * self.save_every) == 0:
val_z = np.random.uniform(-1,
1,
size=(val_size, self.z_dim))
val_out = sess.run(network.generator(self.inputs_z,
reuse=True,
is_training=False),
feed_dict={self.inputs_z: val_z})
image_fn = os.path.join(
self.assets_dir, '{:s}-val-e{:03d}.png'.format(
self.dataset_type, (e // self.d_train_freq + 1)))
utils.validation(val_out,
self.val_block_size,
image_fn,
color_mode='L')
end_time = time.time()
elapsed_time = end_time - start_time
# save losses as image
losses_fn = os.path.join(self.assets_dir,
'{:s}-losses.png'.format(self.dataset_type))
utils.save_losses(losses, ['Discriminator', 'Generator'], elapsed_time,
losses_fn)
return
def __init__(self, args, device):
# parameters
self.epoch = args.epochs
self.batch_size = args.batch_size
self.save_dir = args.save_dir
self.result_dir = args.result_dir
self.dataset = args.dataset
self.log_dir = args.log_dir
self.model_name = "GCGAN"
self.Glayer_num = args.Glayer
self.Dlayer_num = args.Dlayer
self.Ghidden_num = args.Ghidden
self.z_dim = args.z_dim
self.num_worker = args.num_worker
self.device = device
dataset = MovieLensDataset(dataset=self.dataset,
transform=transforms.Compose([ToTensor()]))
dataset_num = len(dataset)
train_num = int(dataset_num * 0.8)
train_dataset, test_dataset = random_split(
dataset, [train_num, dataset_num - train_num])
# load dataset
self.train_loader = DataLoader(train_dataset,
batch_size=self.batch_size,
shuffle=True,
num_workers=self.num_worker)
self.test_loader = DataLoader(test_dataset,
batch_size=len(dataset),
shuffle=True,
num_workers=self.num_worker)
data = dataset[0]['u_perchase']
self.u_feature_num = dataset[0]['u_feature'].shape[0]
self.v_feature_num = dataset[0]['v_feature'].shape[1]
# networks init
self.G = generator(input_dim=self.z_dim,
feature_num=self.u_feature_num,
output_dim=data.shape[0],
layer_num=self.Glayer_num,
hidden_num=self.Ghidden_num).to(self.device)
self.D = discriminator(in_features_u=self.u_feature_num,
num_item=data.shape[0],
in_features_v=self.v_feature_num,
rating=5,
output_dim=1,
layer_num=self.Dlayer_num).to(self.device)
self.G_optimizer = optim.SGD(self.G.parameters(), lr=args.lrG)
self.D_optimizer = optim.SGD(self.D.parameters(), lr=args.lrD)
self.BCE_loss = nn.BCELoss().to(self.device)
self.MSE_loss = nn.MSELoss().to(self.device)
print('---------- Networks architecture -------------')
utils.print_network(self.G)
utils.print_network(self.D)
print('-----------------------------------------------')
def __init__(self, name, dataset_type, gan_loss_type):
# prepare directories
self.assets_dir = './assets/{:s}'.format(name)
self.ckpt_dir = './ckpts/{:s}'.format(name)
self.ckpt_fn = os.path.join(self.ckpt_dir, '{:s}.ckpt'.format(name))
if not os.path.exists(self.assets_dir):
os.makedirs(self.assets_dir)
if not os.path.exists(self.ckpt_dir):
os.makedirs(self.ckpt_dir)
# setup variables
self.dataset_type = dataset_type
# tunable parameters
self.z_dim = 100
self.learning_rate = 1e-4
self.epochs = 30
self.batch_size = 128
self.print_every = 30
self.save_every = 5
self.val_block_size = 10
self.lmbd_gp = 10.0
# start building graphs
tf.reset_default_graph()
# create placeholders
self.running_bs = tf.placeholder(tf.int32, [], name='running_bs')
self.latent_z = tf.placeholder(tf.float32, [None, self.z_dim],
name='latent_z')
self.real_images = tf.placeholder(tf.float32, [None, 28, 28, 1],
name='real_images')
# create generator & discriminator
self.fake_images = network.generator(self.latent_z,
is_training=True,
use_bn=False)
self.d_real_logits, _ = network.discriminator(self.real_images,
is_training=True,
use_bn=False)
self.d_fake_logits, _ = network.discriminator(self.fake_images,
is_training=True,
use_bn=False)
# compute model loss
self.d_loss, self.g_loss = wgan_loss(self.d_real_logits,
self.d_fake_logits)
# add gradient penalty
alpha = tf.random_uniform(shape=[self.running_bs, 1, 1, 1],
minval=-1.0,
maxval=1.0)
interpolates = self.real_images + alpha * (self.fake_images -
self.real_images)
d_interpolates_logits, _ = network.discriminator(interpolates,
is_training=True,
use_bn=False)
gradients = tf.gradients(d_interpolates_logits, [interpolates])[0]
slopes = tf.sqrt(
0.0001 +
tf.reduce_sum(tf.square(gradients), reduction_indices=[1, 2, 3]))
gradient_penalty = tf.reduce_mean(tf.square(slopes - 1.0))
self.d_loss += self.lmbd_gp * gradient_penalty
# prepare optimizers
t_vars = tf.trainable_variables()
d_vars = [
var for var in t_vars if var.name.startswith('discriminator')
]
g_vars = [var for var in t_vars if var.name.startswith('generator')]
# Optimize
optimizer = tf.train.AdamOptimizer(self.learning_rate,
beta1=0.5,
beta2=0.9)
self.d_opt = optimizer.minimize(self.d_loss, var_list=d_vars)
self.g_opt = optimizer.minimize(
self.g_loss,
var_list=g_vars,
global_step=tf.train.get_or_create_global_step())
# prepare saver for generator
self.saver = tf.train.Saver(var_list=g_vars)
return
parser.add_argument('--save_root', required=False, default='results', help='results save path')
parser.add_argument('--inverse_order', type=bool, default=True, help='0: [input, target], 1 - [target, input]')
opt = parser.parse_args()
print(opt)
# data_loader
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
test_loader = util.data_load('data/' + opt.dataset, opt.test_subfolder, transform, batch_size=1, shuffle=False)
if not os.path.isdir(opt.dataset + '_results/test_results'):
os.mkdir(opt.dataset + '_results/test_results')
G = network.generator(opt.ngf)
G.cuda()
G.load_state_dict(torch.load(opt.dataset + '_results/' + opt.dataset + '_generator_param.pkl'))
# network
n = 0
print('test start!')
for x_, _ in test_loader:
if opt.inverse_order:
y_ = x_[:, :, :, :x_.size()[2]]
x_ = x_[:, :, :, x_.size()[2]:]
else:
y_ = x_[:, :, :, x_.size()[2]:]
x_ = x_[:, :, :, :x_.size()[2]]
if x_.size()[2] != opt.input_size:
if not os.path.isdir(os.path.join(args.name + '_results', 'Colorization')):
os.makedirs(os.path.join(args.name + '_results', 'Colorization'))
transform = transforms.Compose([
transforms.Resize((args.input_size, args.input_size)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
train_loader_src = utils.data_load(os.path.join('data', args.src_data), 'train', transform, args.batch_size, shuffle=True, drop_last=True)
train_loader_tgt = utils.data_load(os.path.join('data', args.tgt_data), 'train', transform, args.batch_size, shuffle=True, drop_last=True)
test_loader_src = utils.data_load(os.path.join('data', args.src_data), 'test', transform, 1, shuffle=True, drop_last=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
A2BG = net.generator(args.in_ngc, args.out_ngc, args.ngf)
B2AG = net.generator(args.in_ngc, args.out_ngc, args.ngf)
AD = net.discriminator(args.in_ndc, args.out_ndc, args.ndf)
BD = net.discriminator(args.in_ndc, args.out_ndc, args.ndf)
print('---------- Networks initialized -------------')
utils.print_network(A2BG)
utils.print_network(AD)
print('-----------------------------------------------')
vgg16 = models.vgg16(pretrained=True)
vgg16 = net.VGG(vgg16.features[:23]).to(device)
A2BG.to(device)
B2AG.to(device)
#train_labels_onehot = EncodingOnehot(train_labels, nclasses)
test_labels = LoadLabel(TEST_LABEL, DATA_DIR)
test_labels_onehot = EncodingOnehot(test_labels, nclasses)
train_labels = LoadLabel(DATABASE_LABEL, DATA_DIR)
train_labels_onehot = EncodingOnehot(train_labels, nclasses)
Y = train_labels_onehot
#Sim = CalcSim(train_labels_onehot, train_labels_onehot)
# T_S = np.load('64ex-T_S.npy')
# D_I = np.load('64ex-H_I.npy')
# map_1 = CalcHR.CalcMap(T_S, D_I,test_labels_onehot.numpy(), train_labels_onehot.numpy())
# print(map_1)
# network
G_A = network.generator(opt.input_ngc,
opt.output_ngc,
opt.ngf,
opt.nb,
flag=False)
G_B = network.generator(opt.input_ngc,
opt.output_ngc,
opt.ngf,
opt.nb,
flag=False)
D_A = network.discriminator(opt.input_ndc, opt.bit, opt.ndf)
net = models.resnet18(pretrained=False)
net.fc = nn.Linear(2048, opt.bit)
net_dict = net.state_dict()
pretrain_dict = torch.load('./ex_sketch_64HA_param.pkl')
def train(self):
n_fixed_samples = self.val_block_size * self.val_block_size
fixed_z = np.random.uniform(-1, 1, size=(n_fixed_samples, self.z_dim))
fixed_y = np.zeros(shape=[n_fixed_samples, self.y_dim])
for s in range(n_fixed_samples):
loc = s % self.y_dim
fixed_y[s, loc] = 1
steps = 0
losses = []
start_time = time.time()
with tf.Session() as sess:
# reset tensorflow variables
sess.run(tf.global_variables_initializer())
# start training
for e in range(self.epochs):
for ii in range(self.mnist_loader.train.num_examples //
self.batch_size):
batch_x, batch_y = self.mnist_loader.train.next_batch(
self.batch_size)
# rescale images to -1 ~ 1
batch_x = np.reshape(batch_x, (-1, 28, 28, 1))
batch_x = batch_x * 2.0 - 1.0
# Sample random noise for G
batch_z = np.random.uniform(-1,
1,
size=(self.batch_size,
self.z_dim))
fd = {
self.inputs_x: batch_x,
self.inputs_y: batch_y,
self.inputs_z: batch_z
}
# Run optimizers
_ = sess.run(self.d_opt, feed_dict=fd)
_ = sess.run(self.g_opt, feed_dict=fd)
_ = sess.run(self.ac_opt, feed_dict=fd)
# print losses
if steps % self.print_every == 0:
# At the end of each epoch, get the losses and print them out
train_loss_d = self.d_loss.eval(fd)
train_loss_g = self.g_loss.eval(fd)
train_loss_ac = self.ac_loss.eval(fd)
print(
"Epoch {}/{}...".format(e + 1, self.epochs),
"Discriminator Loss: {:.4f}...".format(
train_loss_d),
"Generator Loss: {:.4f}...".format(train_loss_g),
"Auxilary Classifier Loss: {:.4f}...".format(
train_loss_ac))
losses.append(
(train_loss_d, train_loss_g, train_loss_ac))
steps += 1
# save generation results at every epochs
if e % self.save_every == 0:
val_out = sess.run(network.generator(self.inputs_z,
y=self.inputs_y,
reuse=True,
is_training=False),
feed_dict={
self.inputs_y: fixed_y,
self.inputs_z: fixed_z
})
image_fn = os.path.join(
self.assets_dir,
'{:s}-val-e{:03d}.png'.format(self.dataset_type,
e + 1))
utils.validation(val_out,
self.val_block_size,
image_fn,
color_mode='L')
end_time = time.time()
elapsed_time = end_time - start_time
# save losses as image
losses_fn = os.path.join(self.assets_dir,
'{:s}-losses.png'.format(self.dataset_type))
utils.save_losses(losses, ['Discriminator', 'Generator', 'Auxilary'],
elapsed_time, losses_fn)
return
