def train(self):
# noise for test.
sample_batch = self.loader.get_batch()
print(sample_batch)
self.z_test = sample_batch['encods']
print("0self.z_test")
print(self.z_test)
print("1self.z_test")
if self.use_cuda:
self.z_test = self.z_test.cuda()
self.z_test = Variable(self.z_test, volatile=False)
self.z_test.data.resize_(self.loader.batchsize, self.nz)
for step in range(2, self.max_resl + 1 + 5):
for iter in tqdm(
range(0, (self.trns_tick * 2 + self.stab_tick * 2) *
self.TICK, self.loader.batchsize)):
sample_batch = self.loader.get_batch()
self.globalIter = self.globalIter + 1
self.stack = self.stack + self.loader.batchsize
if self.stack > ceil(len(self.loader.dataset)):
self.epoch = self.epoch + 1
self.stack = int(self.stack %
(ceil(len(self.loader.dataset))))
# reslolution scheduler.
self.resl_scheduler()
# zero gradients.
self.G.zero_grad()
self.D.zero_grad()
# update discriminator.
self.x.data = self.feed_interpolated_input(
sample_batch['image'])
if self.flag_add_noise:
self.x = self.add_noise(self.x)
self.z = sample_batch['encods']
print("2self.z")
print(self.z_test)
print("3self.z")
if self.use_cuda:
self.z = self.z.cuda()
self.z = Variable(self.z, volatile=False)
self.z.data.resize_(self.loader.batchsize, self.nz)
self.x_tilde = self.G(self.z.float())
self.fx = self.D(self.x.float())
self.fx_tilde = self.D(self.x_tilde.detach())
loss_d = self.mse(self.fx, self.real_label) + self.mse(
self.fx_tilde, self.fake_label)
loss_d.backward()
self.opt_d.step()
# update generator.
fx_tilde = self.D(self.x_tilde)
loss_g = self.mse(fx_tilde, self.real_label.detach())
loss_g.backward()
self.opt_g.step()
# logging.
log_msg = ' [E:{0}][T:{1}][{2:6}/{3:6}] errD: {4:.4f} | errG: {5:.4f} | [lr:{11:.5f}][cur:{6:.3f}][resl:{7:4}][{8}][{9:.1f}%][{10:.1f}%]'.format(
self.epoch, self.globalTick, self.stack,
len(self.loader.dataset), loss_d.data[0], loss_g.data[0],
self.resl, int(pow(2, floor(self.resl))), self.phase,
self.complete['gen'], self.complete['dis'], self.lr)
tqdm.write(log_msg)
# save model.
self.snapshot('repo_enco/model')
# save image grid.
if self.globalIter % self.config.save_img_every == 0:
x_test = self.G(self.z_test.float())
os.system('mkdir -p repo_enco/save/grid')
utils.save_image_grid(
x_test.data,
'repo_enco/save/grid/{}_{}_G{}_D{}.jpg'.format(
int(self.globalIter / self.config.save_img_every),
self.phase, self.complete['gen'],
self.complete['dis']))
os.system('mkdir -p repo_enco/save/resl_{}'.format(
int(floor(self.resl))))
utils.save_image_single(
x_test.data,
'repo_enco/save/resl_{}/{}_{}_G{}_D{}.jpg'.format(
int(floor(self.resl)),
int(self.globalIter / self.config.save_img_every),
self.phase, self.complete['gen'],
self.complete['dis']))
# tensorboard visualization.
if self.use_tb:
x_test = self.G(self.z_test)
self.tb.add_scalar('data/loss_g', loss_g.data[0],
self.globalIter)
self.tb.add_scalar('data/loss_d', loss_d.data[0],
self.globalIter)
self.tb.add_scalar('tick/lr', self.lr, self.globalIter)
self.tb.add_scalar('tick/cur_resl',
int(pow(2, floor(self.resl))),
self.globalIter)
self.tb.add_image_grid(
'grid/x_test', 4,
utils.adjust_dyn_range(x_test.data.float(), [-1, 1],
[0, 1]), self.globalIter)
self.tb.add_image_grid(
'grid/x_tilde', 4,
utils.adjust_dyn_range(self.x_tilde.data.float(),
[-1, 1], [0, 1]),
self.globalIter)
self.tb.add_image_grid(
'grid/x_intp', 4,
utils.adjust_dyn_range(self.x.data.float(), [-1, 1],
[0, 1]), self.globalIter)
def train(self):
# noise for test.
self.z_test = torch.FloatTensor(self.loader.batchsize, self.nz)
if self.use_cuda:
self.z_test = self.z_test.cuda()
self.z_test = Variable(self.z_test, volatile=True)
self.z_test.data.resize_(self.loader.batchsize, self.nz).normal_(0.0, 1.0)
if self.use_captions:
test_caps_set = False
self.caps_test = torch.FloatTensor(self.loader.batchsize, self.ncap)
if self.use_cuda:
self.caps_test = self.caps_test.cuda()
self.caps_test = Variable(self.caps_test, volatile=True)
for step in range(2, self.max_resl+1+5):
for iter in tqdm(range(0,(self.trns_tick*2+self.stab_tick*2)*self.TICK, self.loader.batchsize)):
self.globalIter = self.globalIter+1
self.stack = self.stack + self.loader.batchsize
if self.stack > ceil(len(self.loader.dataset)):
self.epoch = self.epoch + 1
self.stack = int(self.stack%(ceil(len(self.loader.dataset))))
# reslolution scheduler.
self.resl_scheduler()
# zero gradients.
self.G.zero_grad()
self.D.zero_grad()
# update discriminator.
if self.use_captions:
batch_imgs, batch_caps = self.loader.get_batch()
if self.use_cuda:
batch_caps = batch_caps.cuda()
self.caps.data = batch_caps
if not test_caps_set:
self.caps_test.data = batch_caps
test_caps_set = True
else:
batch_imgs, _ = self.loader.get_batch()
self.x.data = self.feed_interpolated_input(batch_imgs)
if self.flag_add_noise:
self.x = self.add_noise(self.x)
self.z.data.resize_(self.loader.batchsize, self.nz).normal_(0.0, 1.0)
if not self.use_captions:
self.x_tilde = self.G(self.z)
else:
self.x_tilde = self.G(self.z, self.caps)
if not self.use_captions:
self.fx = self.D(self.x)
self.fx_tilde = self.D(self.x_tilde.detach())
else:
self.fx = self.D(self.x, self.caps)
self.fx_tilde = self.D(self.x_tilde.detach(), self.caps)
if self.gan_type == 'lsgan':
loss_d = self.mse(self.fx, self.real_label) + self.mse(self.fx_tilde, self.fake_label)
elif self.gan_type == 'wgan-gp':
D_real_loss = -torch.mean(self.fx_tilde)
D_fake_loss = torch.mean(self.x_tilde)
if self.use_cuda:
alpha = torch.rand(self.x.size().cuda())
else:
alpha = torch.rand(self.x.size())
x_hat = Variable(alpha * self.x.data + (1- alpha) * self.G.data, requires_grad=True)
pred_hat = self.D(x_hat)
if self.use_cuda:
gradients = grad(outputs=pred_hat, inputs=x_hat, grad_outputs=torch.ones(pred_hat.size()).cuda(),
create_graph=True, retain_graph=True, only_inputs=True)[0]
else:
gradients = grad(outputs=pred_hat, inputs=x_hat, grad_outputs=torch.ones(pred_hat.size()),
create_graph=True, retain_graph=True, only_inputs=True)[0]
gradient_penalty = self.lambda * ((gradients.view(gradients.size()[0], -1).norm(2, 1) - 1) ** 2).mean()
loss_d = D_real_loss + D_fake_loss + gradient_penalty
loss_d.backward()
self.opt_d.step()
# update generator.
if not self.use_captions:
fx_tilde = self.D(self.x_tilde)
else:
fx_tilde = self.D(self.x_tilde, self.caps)
if self.gan_type == 'lsgan':
loss_g = self.mse(fx_tilde, self.real_label.detach())
elif self.gan_type == 'wgan-gp':
loss_g = -torch.mean(fx_tilde)
loss_g.backward()
self.opt_g.step()
# logging.
log_msg = ' [E:{0}][T:{1}][{2:6}/{3:6}] errD: {4:.4f} | errG: {5:.4f} | [lr:{11:.5f}][cur:{6:.3f}][resl:{7:4}][{8}][{9:.1f}%][{10:.1f}%]'.format(self.epoch, self.globalTick, self.stack, len(self.loader.dataset), loss_d.data[0], loss_g.data[0], self.resl, int(pow(2,floor(self.resl))), self.phase, self.complete['gen'], self.complete['dis'], self.lr)
tqdm.write(log_msg)
# save model.
self.snapshot('repo/model')
# save image grid.
if self.globalIter%self.config.save_img_every == 0:
if not self.use_captions:
x_test = self.G(self.z_test)
else:
x_test = self.G(self.z_test, self.caps_test)
os.system('mkdir -p repo/save/grid')
utils.save_image_grid(x_test.data, 'repo/save/grid/{}_{}_G{}_D{}.jpg'.format(int(self.globalIter/self.config.save_img_every), self.phase, self.complete['gen'], self.complete['dis']))
os.system('mkdir -p repo/save/resl_{}'.format(int(floor(self.resl))))
utils.save_image_single(x_test.data, 'repo/save/resl_{}/{}_{}_G{}_D{}.jpg'.format(int(floor(self.resl)),int(self.globalIter/self.config.save_img_every), self.phase, self.complete['gen'], self.complete['dis']))
# tensorboard visualization.
if self.use_tb:
if not self.use_captions:
x_test = self.G(self.z_test)
else:
x_test = self.G(self.z_test, self.caps_test)
self.tb.add_scalar('data/loss_g', loss_g.data[0], self.globalIter)
self.tb.add_scalar('data/loss_d', loss_d.data[0], self.globalIter)
self.tb.add_scalar('tick/lr', self.lr, self.globalIter)
self.tb.add_scalar('tick/cur_resl', int(pow(2,floor(self.resl))), self.globalIter)
self.tb.add_image_grid('grid/x_test', 4, utils.adjust_dyn_range(x_test.data.float(), [-1,1], [0,1]), self.globalIter)
self.tb.add_image_grid('grid/x_tilde', 4, utils.adjust_dyn_range(self.x_tilde.data.float(), [-1,1], [0,1]), self.globalIter)
self.tb.add_image_grid('grid/x_intp', 4, utils.adjust_dyn_range(self.x.data.float(), [-1,1], [0,1]), self.globalIter)
def train(self):
# noise for test
self.z_test = torch.FloatTensor(self.loader.batchsize, self.nz)
if self.use_cuda:
self.z_test = self.z_test.cuda()
self.z_test = Variable(self.z_test, volatile=True)
self.z_test.data.resize_(self.loader.batchsize, self.nz).normal_(0.0, 1.0)
for step in range(0, self.max_resl + 1 + 5):
for iter in tqdm(range(0, (self.trns_tick * 2 + self.stab_tick * 2) * self.TICK, self.loader.batchsize)):
self.global_iter = self.global_iter + 1
self.stack = self.stack + self.loader.batchsize
if self.stack > ceil(len(self.loader.dataset)):
self.epoch = self.epoch + 1
self.stack = int(self.stack % (ceil(len(self.loader.dataset))))
# Resolution scheduler
self.resl_scheduler()
# Zero the gradients
self.G.zero_grad()
self.D.zero_grad()
# Update discriminator
self.x.data = self.feed_interpolated_input(self.loader.get_batch())
if self.flag_add_noise:
self.x = self.add_noise(self.x)
self.z.data.resize_(self.loader.batchsize, self.nz).normal_(0.0, 1.0)
self.x_tilde = self.G(self.z)
self.fx = self.D(self.x)
self.fx_tilde = self.D(self.x_tilde.detach())
real_loss = self.criterion(torch.squeeze(self.fx), self.real_label)
fake_loss = self.criterion(torch.squeeze(self.fx_tilde), self.fake_label)
loss_d = real_loss + fake_loss
# Compute gradients and apply update to parameters
loss_d.backward()
self.opt_d.step()
# Update generator
fx_tilde = self.D(self.x_tilde)
loss_g = self.criterion(torch.squeeze(fx_tilde), self.real_label.detach())
# Compute gradients and apply update to parameters
loss_g.backward()
self.opt_g.step()
# Log information
log_msg = ' [epoch:{0}][T:{1}][{2:6}/{3:6}] errD: {4:.4f} | errG: {5:.4f} | [lr:{11:.5f}][cur:{6:.3f}][resl:{7:4}][{8}][{9:.1f}%][{10:.1f}%]'.format(
self.epoch,
self.global_tick,
self.stack,
len(self.loader.dataset),
loss_d.data[0],
loss_g.data[0],
self.resl,
int(pow(2, floor(self.resl))),
self.phase,
self.complete['gen'],
self.complete['dis'],
self.lr)
tqdm.write(log_msg)
# Save the model
self.snapshot('./repo/model')
# Save the image grid
if self.global_iter % self.config.save_img_every == 0:
x_test = self.G(self.z_test)
os.system('mkdir -p repo/save/grid')
utils.save_image_grid(x_test.data, 'repo/save/grid/{}_{}_G{}_D{}.jpg'.format(int(self.global_iter / self.config.save_img_every), self.phase, self.complete['gen'], self.complete['dis']))
os.system('mkdir -p repo/save/resl_{}'.format(int(floor(self.resl))))
utils.save_image_single(x_test.data, 'repo/save/resl_{}/{}_{}_G{}_D{}.jpg'.format(int(floor(self.resl)), int(self.global_iter / self.config.save_img_every), self.phase, self.complete['gen'], self.complete['dis']))
# Tensorboard visualization
if self.use_tb:
x_test = self.G(self.z_test)
self.tb.add_scalar('data/loss_g', loss_g.data[0], self.global_iter)
self.tb.add_scalar('data/loss_d', loss_d.data[0], self.global_iter)
self.tb.add_scalar('tick/lr', self.lr, self.global_iter)
self.tb.add_scalar('tick/cur_resl', int(pow(2,floor(self.resl))), self.global_iter)
self.tb.add_image_grid('grid/x_test', 4, utils.adjust_dyn_range(x_test.data.float(), [-1, 1], [0, 1]), self.global_iter)
self.tb.add_image_grid('grid/x_tilde', 4, utils.adjust_dyn_range(self.x_tilde.data.float(), [-1, 1], [0, 1]), self.global_iter)
self.tb.add_image_grid('grid/x_intp', 4, utils.adjust_dyn_range(self.x.data.float(), [-1, 1], [0, 1]), self.global_iter)
def train(self):
# noise for test.
self.z_test = torch.FloatTensor(self.loader.batch_size, self.nz)
print(self.z_test.size())
if self.use_cuda:
self.z_test = self.z_test.cuda()
self.z_test = Variable(self.z_test, volatile=True)
self.z_test.data.resize_(self.loader.batch_size,
self.nz).normal_(0.0, 1.0)
for step in range(2, self.max_resl + 1 + 5):
for iter in tqdm(
range(0, (self.trns_tick * 2 + self.stab_tick * 2) *
self.TICK, self.loader.batch_size)):
self.globalIter = self.globalIter + 1
self.stack = self.stack + self.loader.batch_size
if self.stack > ceil(len(self.loader.dataset)):
self.epoch = self.epoch + 1
self.stack = int(self.stack %
(ceil(len(self.loader.dataset))))
# reslolution scheduler.
self.resl_scheduler()
# zero gradients.
self.G.zero_grad()
self.D.zero_grad()
# update discriminator.
self.x.data = self.feed_interpolated_input(
self.loader.get_batch())
if self.flag_add_noise:
self.x = self.add_noise(self.x)
self.z.data.resize_(self.loader.batch_size,
self.nz).normal_(0.0, 1.0)
self.x_tilde = self.G(self.z)
self.fx = self.D(self.x)
self.fx_tilde = self.D(self.x_tilde.detach())
loss_d = self.mse(self.fx.squeeze(),
self.real_label) + self.mse(
self.fx_tilde, self.fake_label)
loss_d.backward()
self.opt_d.step()
# update generator.
fx_tilde = self.D(self.x_tilde)
identity_loss = 0
for image, latent in self.identity_loader:
image = Variable(image.type(torch.Tensor))
latent = Variable(latent)
gen_image = self.G(latent)
identity_loss = self.l1(gen_image, image) * 0.5
break
loss_g = self.mse(fx_tilde.squeeze(),
self.real_label.detach()) + identity_loss
loss_g.backward()
self.opt_g.step()
# logging.
log_msg = ' [E:{0}][T:{1}][{2:6}/{3:6}] errD: {4:.4f} | errG: {5:.4f} | [lr:{11:.5f}][cur:{6:.3f}][resl:{7:4}][{8}][{9:.1f}%][{10:.1f}%]'.format(
self.epoch, self.globalTick, self.stack,
len(self.loader.dataset), loss_d.item(), loss_g.item(),
self.resl, int(pow(2, floor(self.resl))), self.phase,
self.complete['gen'], self.complete['dis'], self.lr)
tqdm.write(log_msg)
# save image grid.
if self.globalIter % self.config.save_img_every == 0:
with torch.no_grad():
x_test = self.G(self.z_test)
for _, latent in self.identity_loader:
self.save_image(self.G(latent), 'identity/')
self.save_image(x_test, 'grid/')
self.save_image(x_test,
'resl_{}/'.format(int(floor(self.resl))))
# tensorboard visualization.
if self.use_tb:
with torch.no_grad():
x_test = self.G(self.z_test)
self.tb.add_scalar('image300/loss_g', loss_g.item(),
self.globalIter)
self.tb.add_scalar('image300/loss_d', loss_d.item(),
self.globalIter)
self.tb.add_scalar('tick/lr', self.lr, self.globalIter)
self.tb.add_scalar('tick/cur_resl',
int(pow(2, floor(self.resl))),
self.globalIter)
# '''IMAGE GRID
self.tb.add_image_grid(
'grid/x_test', 4,
utils.adjust_dyn_range(x_test.data.float(), [-1, 1],
[0, 1]), self.globalIter)
self.tb.add_image_grid(
'grid/x_tilde', 4,
utils.adjust_dyn_range(self.x_tilde.data.float(),
[-1, 1], [0, 1]),
self.globalIter)
self.tb.add_image_grid(
'grid/x_intp', 4,
utils.adjust_dyn_range(self.x.data.float(), [-1, 1],
[0, 1]), self.globalIter)