def update_core(self):
xp = self.gen.xp
self.gen.cleargrads()
self.dis.cleargrads()
if self.config.generator_architecture == "stylegan":
opt_g_m = self.get_optimizer('map')
opt_g_g = self.get_optimizer('gen')
opt_d = self.get_optimizer('dis')
# z: latent | x: data | y: dis output
# *_real/*_fake/*_pertubed: Variable
# *_data: just data (xp array)
stage = self.stage # Need to retrive the value since next statement may change state (at the stage boundary)
batch = self.get_iterator('main').next()
batch_size = len(batch)
# lr_scale = get_lr_scale_factor(self.total_gpu, stage)
x_real_data = self.get_x_real_data(batch, batch_size)
z_fake_data = xp.concatenate([self.get_z_fake_data(batch_size // 2)] * 2) # repeat same z
if isinstance(chainer.global_config.dtype, chainer._Mixed16):
x_real_data = x_real_data.astype("float16")
z_fake_data = z_fake_data.astype("float16")
x_real = Variable(x_real_data)
# Image.fromarray(convert_batch_images(x_real.data.get(), 4, 4)).save('no_downsized.png')
x_real = downsize_real(x_real, stage)
x_real = Variable(x_real.data)
# Image.fromarray(convert_batch_images(x_real.data.get(), 4, 4)).save('downsized.png')
image_size = x_real.shape[2]
x_fake = self.gen(z_fake_data, stage)
y_fake = self.dis(x_fake[:, :3], stage=stage)
loss_gen = loss_func_dcgan_gen(y_fake) # * lr_scale
assert not xp.isnan(loss_gen.data)
chainer.report({'loss_adv': loss_gen}, opt_g_g.target)
if chainer.global_config.debug:
g = c.build_computational_graph(loss_gen)
with open('out_loss_gen', 'w') as o:
o.write(g.dump())
# assert not xp.isnan(loss_dsgan.data)
# with chainer.using_config('debug', True):
loss_gen.backward()
if self.config.generator_architecture == "stylegan":
opt_g_m.update()
opt_g_g.update()
del loss_gen
self.dis.cleargrads()
# keep smoothed generator if instructed to do so.
if self.smoothed_gen is not None:
# layers_in_use = self.gen.get_layers_in_use(stage=stage)
soft_copy_param(self.smoothed_gen, self.gen, 1.0 - self.smoothing)
y_fake = self.dis(Variable(x_fake.array[:, :3]), stage=stage)
y_real = self.dis(x_real, stage=stage)
loss_dis = loss_func_dcgan_dis(y_fake, y_real)
# loss_reg_camera_param = calc_distance(est_camera_param_real, xp.array(thetas)) / 10
if not self.dis.sn and self.lambda_gp > 0:
# y_perturbed = self.dis(x_perturbed, stage=stage)
grad_x_perturbed, = chainer.grad([y_real], [x_real], enable_double_backprop=True)
grad_l2 = F.sqrt(F.sum(grad_x_perturbed ** 2, axis=(1, 2, 3)))
loss_gp = self.lambda_gp * loss_l2(grad_l2, 0.0)
chainer.report({'loss_gp': loss_gp}, self.dis)
loss_dis = loss_dis + loss_gp # * lr_scale
assert not xp.isnan(loss_dis.data)
chainer.report({'loss_adv': loss_dis}, self.dis)
loss_dis.backward()
opt_d.update()
chainer.reporter.report({'stage': stage})
chainer.reporter.report({'batch_size': batch_size})
chainer.reporter.report({'image_size': image_size})
def update_core(self):
xp = self.gen.xp
self.map.cleargrads()
self.gen.cleargrads()
self.dis.cleargrads()
opt_g_m = self.get_optimizer('map')
opt_g_g = self.get_optimizer('gen')
opt_d = self.get_optimizer('dis')
# z: latent | x: data | y: dis output
# *_real/*_fake/*_pertubed: Variable
# *_data: just data (xp array)
stage = self.stage # Need to retrive the value since next statement may change state (at the stage boundary)
batch = self.get_iterator('main').next()
batch_size = len(batch)
lr_scale = get_lr_scale_factor(self.total_gpu, stage)
x_real_data = self.get_x_real_data(batch, batch_size)
z_fake_data = self.get_z_fake_data(batch_size)
x_real = Variable(x_real_data)
# Image.fromarray(convert_batch_images(x_real.data.get(), 4, 4)).save('no_downsized.png')
x_real = downsize_real(x_real, stage)
x_real = Variable(x_real.data)
# Image.fromarray(convert_batch_images(x_real.data.get(), 4, 4)).save('downsized.png')
image_size = x_real.shape[2]
z_fake = Variable(z_fake_data)
w_fake = self.map(z_fake)
if self.style_mixing_rate > 0 and np.random.rand() < self.style_mixing_rate:
z_fake2 = Variable(self.get_z_fake_data(batch_size))
w_fake2 = self.map(z_fake2)
x_fake = self.gen(w_fake, stage=stage, w2=w_fake2)
else:
x_fake = self.gen(w_fake, stage=stage)
y_fake = self.dis(x_fake, stage=stage)
loss_gen = loss_func_dcgan_gen(y_fake) * lr_scale
if chainer.global_config.debug:
g = c.build_computational_graph(loss_gen)
with open('out_loss_gen', 'w') as o:
o.write(g.dump())
assert not xp.isnan(loss_gen.data)
chainer.report({'loss_adv': loss_gen}, self.gen)
loss_gen.backward()
opt_g_m.update()
opt_g_g.update()
# keep smoothed generator if instructed to do so.
if self.smoothed_gen is not None:
# layers_in_use = self.gen.get_layers_in_use(stage=stage)
soft_copy_param(self.smoothed_gen, self.gen, 1.0 - self.smoothing)
soft_copy_param(self.smoothed_map, self.map, 1.0 - self.smoothing)
z_fake_data = self.get_z_fake_data(batch_size)
z_fake = Variable(z_fake_data)
with chainer.using_config('enable_backprop', False):
w_fake = self.map(z_fake)
if self.style_mixing_rate > 0 and np.random.rand() < self.style_mixing_rate:
z_fake2 = Variable(self.get_z_fake_data(batch_size))
w_fake2 = self.map(z_fake2)
x_fake = self.gen(w_fake, stage=stage, w2=w_fake2)
else:
x_fake = self.gen(w_fake, stage=stage)
x_fake.unchain_backward()
y_fake = self.dis(x_fake, stage=stage)
y_real = self.dis(x_real, stage=stage)
loss_adv = loss_func_dcgan_dis(y_fake, y_real)
if self.lambda_gp > 0:
x_perturbed = x_real
y_perturbed = y_real
# y_perturbed = self.dis(x_perturbed, stage=stage)
grad_x_perturbed, = chainer.grad([y_perturbed], [x_perturbed], enable_double_backprop=True)
grad_l2 = F.sqrt(F.sum(grad_x_perturbed ** 2, axis=(1, 2, 3)))
loss_gp = self.lambda_gp * loss_l2(grad_l2, 0.0)
chainer.report({'loss_gp': loss_gp}, self.dis)
else:
loss_gp = 0.
loss_dis = (loss_adv + loss_gp) * lr_scale
assert not xp.isnan(loss_dis.data)
chainer.report({'loss_adv': loss_adv}, self.dis)
self.dis.cleargrads()
loss_dis.backward()
opt_d.update()
chainer.reporter.report({'stage': stage})
chainer.reporter.report({'batch_size': batch_size})
chainer.reporter.report({'image_size': image_size})
def update_core(self):
xp = self.gen.xp
use_rotate = True if self.iteration > self.config.start_rotation else False
self.gen.cleargrads()
self.dis.cleargrads()
if self.bigan:
self.enc.cleargrads()
if self.config.generator_architecture == "stylegan":
opt_g_m = self.get_optimizer('map')
opt_g_g = self.get_optimizer('gen')
opt_d = self.get_optimizer('dis')
# z: latent | x: data | y: dis output
# *_real/*_fake/*_pertubed: Variable
# *_data: just data (xp array)
stage = self.stage # Need to retrive the value since next statement may change state (at the stage boundary)
batch = self.get_iterator('main').next()
batch_size = len(batch)
# lr_scale = get_lr_scale_factor(self.total_gpu, stage)
x_real_data = self.get_x_real_data(batch, batch_size)
z_fake_data = xp.concatenate([self.get_z_fake_data(batch_size // 2)] * 2) # repeat same z
if isinstance(chainer.global_config.dtype, chainer._Mixed16):
x_real_data = x_real_data.astype("float16")
z_fake_data = z_fake_data.astype("float16")
# theta->6 DOF
thetas = self.prior.sample(batch_size)
# thetas = Variable(xp.array(thetas))
# # theta -> camera matrix
# random_camera_matrices = xp.array(get_camera_matries(thetas, order=(0, 1, 2)), dtype="float32")
# thetas = F.concat([F.cos(thetas[:, :3]), F.sin(thetas[:, :3]), thetas[:, 3:]], axis=1)
# theta -> camera matrix
thetas_ = xp.array(thetas)
random_camera_matrices = xp.array(get_camera_matries(thetas))
thetas = F.concat([F.cos(thetas_[:, :3]), F.sin(thetas_[:, :3]),
thetas_[:, 3:]], axis=1)
x_real = Variable(x_real_data)
# Image.fromarray(convert_batch_images(x_real.data.get(), 4, 4)).save('no_downsized.png')
x_real = downsize_real(x_real, stage)
x_real = Variable(x_real.data)
# Image.fromarray(convert_batch_images(x_real.data.get(), 4, 4)).save('downsized.png')
image_size = x_real.shape[2]
x_fake = self.gen(z_fake_data, stage, thetas)
if self.bigan:
# bigan is not supported now
assert False, "bigan is not supported"
else:
y_fake, feat = self.dis(x_fake[:, :3], stage=stage, return_hidden=True)
loss_gen = loss_func_dcgan_gen(y_fake) # * lr_scale
assert not xp.isnan(loss_gen.data)
chainer.report({'loss_adv': loss_gen}, opt_g_g.target)
if use_rotate:
loss_rotate, warped_zp = self.loss_func_rotate(x_fake[:batch_size // 2],
random_camera_matrices[:batch_size // 2],
x_fake[batch_size // 2:],
random_camera_matrices[batch_size // 2:],
self.iteration >= self.config.start_occlusion_aware)
if self.config.rotate_feature:
downsample_rate = x_real.shape[2] // feat.shape[2]
depth = F.average_pooling_2d(x_real[:, -1:], downsample_rate, downsample_rate, 0)
feat = F.concat([feat, depth], axis=1)
loss_rotate_feature, _ = self.loss_func_rotate_feature(feat[:batch_size // 2],
random_camera_matrices[:batch_size // 2],
feat[batch_size // 2:],
random_camera_matrices[batch_size // 2:],
self.iteration >= self.config.start_occlusion_aware)
loss_rotate += loss_rotate_feature
# loss_rotate *= 10
if self.config.lambda_depth > 0:
loss_rotate += F.mean(F.relu(self.config.depth_min - x_fake[:, -1]) ** 2) * \
self.config.lambda_depth # make depth larger
assert not xp.isnan(loss_rotate.data)
chainer.report({'loss_rotate': loss_rotate}, opt_g_g.target)
lambda_rotate = self.config.lambda_rotate if self.config.lambda_rotate else 2
lambda_rotate = lambda_rotate if image_size <= 128 else lambda_rotate * 2
loss_gen += loss_rotate * lambda_rotate
if self.config.use_occupancy_net_loss:
loss_occupancy = self.loss_func_rotate.occupancy_net_loss(self.gen.occupancy, x_fake[:, -1:],
random_camera_matrices, z_fake_data.squeeze())
chainer.report({'loss_occupancy': loss_occupancy}, opt_g_g.target)
loss_gen += loss_occupancy * self.config.lambda_occupancy
if self.config.optical_flow:
assert False, "optical flow loss is not supported"
# loss_rotate += - 0.2 * F.log(
# F.mean((F.mean(1 / x_fake[:, -1].reshape(batch_size, -1) ** 2, axis=1) -
# F.mean(1 / x_fake[:, -1].reshape(batch_size, -1), axis=1) ** 2) + 1e-3))
# loss_depth = self.loss_smooth_depth(x_fake[:, -1:]) * 20
# loss_dsgan = loss_func_dsgan(x_fake, z_fake, theta) # Diversity sensitive gan in ICLR2019
if chainer.global_config.debug:
g = c.build_computational_graph(loss_gen)
with open('out_loss_gen', 'w') as o:
o.write(g.dump())
# assert not xp.isnan(loss_dsgan.data)
# with chainer.using_config('debug', True):
loss_gen.backward()
# loss_depth.backward()
# loss_dsgan.backward()
if self.config.generator_architecture == "stylegan":
opt_g_m.update()
opt_g_g.update()
del loss_gen
self.dis.cleargrads()
# keep smoothed generator if instructed to do so.
if self.smoothed_gen is not None:
# layers_in_use = self.gen.get_layers_in_use(stage=stage)
soft_copy_param(self.smoothed_gen, self.gen, 1.0 - self.smoothing)
# with chainer.using_config('enable_backprop', False):
if self.bigan:
assert False, "bigan is not supported"
else:
v_x_fake = Variable(x_fake.array[:, :3])
y_fake, feat = self.dis(v_x_fake, stage=stage, return_hidden=True)
y_real = self.dis(x_real, stage=stage)
loss_dis = loss_func_dcgan_dis(y_fake, y_real)
# loss_reg_camera_param = calc_distance(est_camera_param_real, xp.array(thetas)) / 10
if not self.dis.sn and self.lambda_gp > 0:
# y_perturbed = self.dis(x_perturbed, stage=stage)
grad_x_perturbed, = chainer.grad([y_real], [x_real], enable_double_backprop=True)
grad_l2 = F.sqrt(F.sum(grad_x_perturbed ** 2, axis=(1, 2, 3)))
loss_gp = self.lambda_gp * loss_l2(grad_l2, 0.0)
chainer.report({'loss_gp': loss_gp}, self.dis)
loss_dis = loss_dis + loss_gp # * lr_scale
if use_rotate and self.config.rotate_feature:
downsample_rate = x_real.shape[2] // feat.shape[2]
depth = F.average_pooling_2d(x_real[:, -1:], downsample_rate, downsample_rate, 0)
feat = F.concat([feat, depth], axis=1)
loss_rotate_feature, _ = self.loss_func_rotate_feature(feat[:batch_size // 2],
random_camera_matrices[:batch_size // 2],
feat[batch_size // 2:],
random_camera_matrices[batch_size // 2:],
self.iteration >= self.config.start_occlusion_aware)
loss_dis -= loss_rotate_feature
if not self.dis.sn and self.lambda_gp > 0:
grad_x_perturbed, = chainer.grad([feat], [v_x_fake], enable_double_backprop=True)
grad_l2 = F.sqrt(F.sum(grad_x_perturbed ** 2, axis=(1, 2, 3)))
loss_gp = self.lambda_gp * loss_l2(grad_l2, 0.0)
loss_dis += loss_gp
assert not xp.isnan(loss_dis.data)
chainer.report({'loss_adv': loss_dis}, self.dis)
loss_dis.backward()
opt_d.update()
chainer.reporter.report({'stage': stage})
chainer.reporter.report({'batch_size': batch_size})
chainer.reporter.report({'image_size': image_size})
def update_core(self):
xp = self.gen.xp
use_rotate = True if self.iteration > self.config.start_rotation else False
self.gen.cleargrads()
self.gen.mapping.cleargrads()
self.dis.cleargrads()
opt_g_m = self.get_optimizer('map')
opt_g_g = self.get_optimizer('gen')
opt_d = self.get_optimizer('dis')
# z: latent | x: data | y: dis output
# *_real/*_fake/*_pertubed: Variable
# *_data: just data (xp array)
stage = self.stage # Need to retrive the value since next statement may change state (at the stage boundary)
batch = self.get_iterator('main').next()
batch_size = len(batch)
# lr_scale = get_lr_scale_factor(self.total_gpu, stage)
x_real_data = self.get_x_real_data(batch, batch_size)
z_fake_data = xp.tile(self.get_z_fake_data(batch_size // 2),
(2, 1, 1, 1, 1)) # repeat same z
z_fake_data2 = xp.tile(self.get_z_fake_data(batch_size // 2),
(2, 1, 1, 1, 1)) # repeat same z
if isinstance(chainer.global_config.dtype, chainer._Mixed16):
x_real_data = x_real_data.astype("float16")
z_fake_data = z_fake_data.astype("float16")
z_fake_data2 = z_fake_data.astype("float16")
# theta->6 DOF
thetas = self.prior.sample(batch_size)
# theta -> camera matrix
random_camera_matrices = xp.array(get_camera_matries(thetas),
dtype="float32")
thetas = xp.array(
np.concatenate(
[np.cos(thetas[:, :3]),
np.sin(thetas[:, :3]), thetas[:, 3:]],
axis=1))
x_real = Variable(x_real_data)
x_real = downsize_real(x_real, IMG_SIZE)
x_real = Variable(x_real.data)
image_size = x_real.shape[2]
x_fake = self.gen(z_fake_data,
stage,
random_camera_matrices,
z2=z_fake_data2,
theta=thetas)
y_fake = self.dis(x_fake[:, :3], stage=stage)
loss_gen = loss_func_dcgan_gen(
y_fake, self.config.focal_loss_gamma) # * lr_scale
chainer.report({'loss_adv': loss_gen}, self.gen)
assert not xp.isnan(loss_gen.data)
if use_rotate:
if self.config.background_generator:
loss_rotate_fore, _ = self.loss_func_rotate(
x_fake[:batch_size // 2],
random_camera_matrices[:batch_size // 2],
x_fake[batch_size // 2:],
random_camera_matrices[batch_size // 2:],
max_depth=3)
virtual_camera_matrices = random_camera_matrices.copy()
virtual_camera_matrices[:, :3, 3] = 0
loss_rotate_back, _ = self.loss_func_rotate(
x_fake[:batch_size // 2],
virtual_camera_matrices[:batch_size // 2],
x_fake[batch_size // 2:],
virtual_camera_matrices[batch_size // 2:],
min_depth=3)
loss_rotate = loss_rotate_fore + loss_rotate_back
else:
loss_rotate, _ = self.loss_func_rotate(
x_fake[:batch_size // 2],
random_camera_matrices[:batch_size // 2],
x_fake[batch_size // 2:],
random_camera_matrices[batch_size // 2:])
loss_rotate += F.mean(F.relu(self.config.depth_min - x_fake[:, -1]) ** 2) * \
self.config.lambda_depth # make depth larger
chainer.report({'loss_rotate': loss_rotate}, self.gen)
assert not xp.isnan(loss_rotate.data)
lambda_loss_rotate = self.config.lambda_loss_rotate if self.config.lambda_loss_rotatec else 0.3
loss_gen = loss_gen + loss_rotate * lambda_loss_rotate
if chainer.global_config.debug:
g = c.build_computational_graph(loss_gen)
with open('out_loss_gen', 'w') as o:
o.write(g.dump())
# assert not xp.isnan(loss_dsgan.data)
loss_gen.backward()
opt_g_m.update()
opt_g_g.update()
del loss_gen, y_fake, x_fake
self.dis.cleargrads()
# keep smoothed generator if instructed to do so.
if self.smoothed_gen is not None:
# layers_in_use = self.gen.get_layers_in_use(stage=stage)
soft_copy_param(self.smoothed_gen, self.gen, 1.0 - self.smoothing)
z_fake_data = self.get_z_fake_data(batch_size)
z_fake_data2 = self.get_z_fake_data(batch_size)
if isinstance(chainer.global_config.dtype, chainer._Mixed16):
z_fake_data = z_fake_data.astype("float16")
z_fake_data2 = z_fake_data.astype("float16")
# with chainer.using_config('enable_backprop', False):
x_fake = self.gen(z_fake_data,
stage,
random_camera_matrices,
z2=z_fake_data2,
theta=thetas)
x_fake.unchain_backward()
y_fake = self.dis(x_fake[:, :3], stage=stage)
y_real = self.dis(x_real, stage=stage)
loss_adv = loss_func_dcgan_dis(y_fake, y_real)
if not self.dis.sn and self.lambda_gp > 0:
x_perturbed = x_real
y_perturbed = y_real
# y_perturbed = self.dis(x_perturbed, stage=stage)
grad_x_perturbed, = chainer.grad([y_perturbed], [x_perturbed],
enable_double_backprop=True)
grad_l2 = F.sqrt(F.sum(grad_x_perturbed**2, axis=(1, 2, 3)))
loss_gp = self.lambda_gp * loss_l2(grad_l2, 0.0)
chainer.report({'loss_gp': loss_gp}, self.dis)
else:
loss_gp = 0.
loss_dis = (loss_adv + loss_gp) # * lr_scale
assert not xp.isnan(loss_dis.data)
chainer.report({'loss_adv': loss_adv}, self.dis)
loss_dis.backward()
opt_d.update()
chainer.reporter.report({'batch_size': batch_size})
chainer.reporter.report({'image_size': image_size})