def __init__(self, args):
args.deterministic = True
if args.dataset in ["mnist", "dsprites"]:
content_encoder = models.content_encoder(
args.g_dims, nc=args.num_channels).cuda()
position_encoder = models.pose_encoder(
args.z_dims,
nc=args.num_channels,
normalize=args.normalize_position).cuda()
else:
content_encoder = vgg_64.encoder(args.g_dims,
nc=args.num_channels).cuda()
position_encoder = resnet_64.pose_encoder(
args.z_dims, nc=args.num_channels).cuda()
if args.dataset == "mpi3d_real":
decoder = vgg_64.drnet_decoder(args.g_dims,
args.z_dims,
nc=args.num_channels).cuda()
else:
decoder = models.decoder(args.g_dims,
args.z_dims,
nc=args.num_channels,
skips=args.skips).cuda()
self.content_frames = 1
if args.content_lstm:
content_encoder = models.content_encoder_lstm(
args.g_dims, content_encoder, args.batch_size)
self.content_frames = args.input_frames
discriminator = models.scene_discriminator(args.z_dims).cuda()
nets = {
"content_encoder": content_encoder,
"position_encoder": position_encoder,
"decoder": decoder,
"discriminator": discriminator,
}
self.encoder_decoder_parameters = itertools.chain(*[
content_encoder.parameters(),
position_encoder.parameters(),
decoder.parameters(),
])
encoder_decoder_optim = torch.optim.Adam(
self.encoder_decoder_parameters,
lr=args.lr,
betas=(args.beta1, 0.999))
discriminator_optim = torch.optim.Adam(discriminator.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
optims = {
"encoder_decoder_optim": encoder_decoder_optim,
"discriminator_optim": discriminator_optim,
}
super().__init__(nets, optims, args)
def __init__(self, args):
args.deterministic = True
encoder_checkpoint = torch.load(args.encoder_checkpoint)
if args.dataset in ["mnist", "dsprites"]:
Ec = models.content_encoder(args.g_dims,
nc=args.num_channels).cuda()
Ep = models.pose_encoder(args.z_dims, nc=args.num_channels).cuda()
else:
Ec = vgg_64.encoder(args.g_dims, nc=args.num_channels).cuda()
Ep = resnet_64.pose_encoder(args.z_dims,
nc=args.num_channels).cuda()
if args.dataset == "mpi3d_real":
D = vgg_64.drnet_decoder(args.g_dims,
args.z_dims,
nc=args.num_channels).cuda()
else:
D = models.decoder(args.g_dims,
args.z_dims,
nc=args.num_channels,
skips=args.skips).cuda()
Ep.load_state_dict(encoder_checkpoint["position_encoder"])
Ec.load_state_dict(encoder_checkpoint["content_encoder"])
D.load_state_dict(encoder_checkpoint["decoder"])
self.Ep = nn.DataParallel(Ep)
self.Ec = nn.DataParallel(Ec)
self.D = nn.DataParallel(D)
self.Ep.train()
self.Ec.train()
self.D.train()
lstm_model = lstm(args.g_dims + args.z_dims, args.z_dims,
args.rnn_size, args.rnn_layers,
args.batch_size).cuda()
nets = {"lstm": lstm_model}
lstm_optim = torch.optim.Adam(lstm_model.parameters(),
lr=args.lr,
betas=(args.beta1, 0.999))
optims = {"lstm_optim": lstm_optim}
super().__init__(nets, optims, args)
def generate(args):
# Load model
nn.load_parameters(args.model_load_path)
# Context
extension_module = "cudnn"
ctx = get_extension_context(extension_module, type_config=args.type_config)
nn.set_default_context(ctx)
# Input
b, c, h, w = 1, 3, args.image_size, args.image_size
x_real_a = nn.Variable([b, c, h, w])
x_real_b = nn.Variable([b, c, h, w])
one = nn.Variable.from_numpy_array(np.ones((1, 1, 1, 1)) * 0.5)
# Model
maps = args.maps
# content/style (domain A)
x_content_a = content_encoder(x_real_a, maps, name="content-encoder-a")
x_style_a = style_encoder(x_real_a, maps, name="style-encoder-a")
# content/style (domain B)
x_content_b = content_encoder(x_real_b, maps, name="content-encoder-b")
x_style_b = style_encoder(x_real_b, maps, name="style-encoder-b")
# generate over domains and reconstruction of content and style (domain A)
z_style_a = F.randn(
shape=x_style_a.shape) if not args.example_guided else x_style_a
z_style_a = z_style_a.apply(persistent=True)
x_fake_a = decoder(x_content_b, z_style_a, name="decoder-a")
# generate over domains and reconstruction of content and style (domain B)
z_style_b = F.randn(
shape=x_style_b.shape) if not args.example_guided else x_style_b
z_style_b = z_style_b.apply(persistent=True)
x_fake_b = decoder(x_content_a, z_style_b, name="decoder-b")
# Monitor
suffix = "Stochastic" if not args.example_guided else "Example-guided"
monitor = Monitor(args.monitor_path)
monitor_image_a = MonitorImage("Fake Image B to A {} Valid".format(suffix),
monitor,
interval=1)
monitor_image_b = MonitorImage("Fake Image A to B {} Valid".format(suffix),
monitor,
interval=1)
# DataIterator
di_a = munit_data_iterator(args.img_path_a, args.batch_size)
di_b = munit_data_iterator(args.img_path_b, args.batch_size)
# Generate all
# generate (A -> B)
if args.example_guided:
x_real_b.d = di_b.next()[0]
for i in range(di_a.size):
x_real_a.d = di_a.next()[0]
images = []
images.append(x_real_a.d.copy())
for _ in range(args.num_repeats):
x_fake_b.forward(clear_buffer=True)
images.append(x_fake_b.d.copy())
monitor_image_b.add(i, np.concatenate(images, axis=3))
# generate (B -> A)
if args.example_guided:
x_real_a.d = di_a.next()[0]
for i in range(di_b.size):
x_real_b.d = di_b.next()[0]
images = []
images.append(x_real_b.d.copy())
for _ in range(args.num_repeats):
x_fake_a.forward(clear_buffer=True)
images.append(x_fake_a.d.copy())
monitor_image_a.add(i, np.concatenate(images, axis=3))
shuffle=True,
num_workers=5,
drop_last=True)
def get_data_batch():
while True:
for seq in data_loader:
seq[1].transpose_(2, 3).transpose_(1, 2)
yield seq
data_generator = get_data_batch()
if args.dataset == "dsprites":
Ec = models.content_encoder(args.g_dims, nc=args.num_channels).cuda()
Ep = models.pose_encoder(args.z_dims, nc=args.num_channels).cuda()
elif args.dataset == "mpi3d_real":
Ec = vgg_64.encoder(args.g_dims, nc=args.num_channels).cuda()
Ep = resnet_64.pose_encoder(args.z_dims, nc=args.num_channels).cuda()
checkpoint = torch.load(args.checkpoint)
Ec.load_state_dict(checkpoint["content_encoder"])
Ep.load_state_dict(checkpoint["position_encoder"])
Ec.eval()
Ep.eval()
latent_c = None
latent_p = None
factors_p = None
def interpolate(args):
# Load model
nn.load_parameters(args.model_load_path)
# Context
extension_module = "cudnn"
ctx = get_extension_context(extension_module, type_config=args.type_config)
nn.set_default_context(ctx)
# Input
b, c, h, w = 1, 3, args.image_size, args.image_size
x_real_a = nn.Variable([b, c, h, w])
x_real_b = nn.Variable([b, c, h, w])
one = nn.Variable.from_numpy_array(np.ones((1, 1, 1, 1)) * 0.5)
# Model
maps = args.maps
# content/style (domain A)
x_content_a = content_encoder(x_real_a, maps, name="content-encoder-a")
x_style_a = style_encoder(x_real_a, maps, name="style-encoder-a")
# content/style (domain B)
x_content_b = content_encoder(x_real_b, maps, name="content-encoder-b")
x_style_b = style_encoder(x_real_b, maps, name="style-encoder-b")
# generate over domains and reconstruction of content and style (domain A)
z_style_a = nn.Variable(
x_style_a.shape) if not args.example_guided else x_style_a
z_style_a = z_style_a.apply(persistent=True)
x_fake_a = decoder(x_content_b, z_style_a, name="decoder-a")
# generate over domains and reconstruction of content and style (domain B)
z_style_b = nn.Variable(
x_style_b.shape) if not args.example_guided else x_style_b
z_style_b = z_style_b.apply(persistent=True)
x_fake_b = decoder(x_content_a, z_style_b, name="decoder-b")
# Monitor
def file_names(path):
return path.split("/")[-1].rstrip("_AB.jpg")
suffix = "Stochastic" if not args.example_guided else "Example-guided"
monitor = Monitor(args.monitor_path)
monitor_image_tile_a = MonitorImageTile(
"Fake Image Tile {} B to A {} Interpolation".format(
"-".join([file_names(path) for path in args.img_files_b]), suffix),
monitor,
interval=1,
num_images=len(args.img_files_b))
monitor_image_tile_b = MonitorImageTile(
"Fake Image Tile {} A to B {} Interpolation".format(
"-".join([file_names(path) for path in args.img_files_a]), suffix),
monitor,
interval=1,
num_images=len(args.img_files_a))
# DataIterator
di_a = munit_data_iterator(args.img_files_a, b, shuffle=False)
di_b = munit_data_iterator(args.img_files_b, b, shuffle=False)
rng = np.random.RandomState(args.seed)
# Interpolate (A -> B)
z_data_0 = [rng.randn(*z_style_a.shape) for j in range(di_a.size)]
z_data_1 = [rng.randn(*z_style_a.shape) for j in range(di_a.size)]
for i in range(args.num_repeats):
r = 1.0 * i / args.num_repeats
images = []
for j in range(di_a.size):
x_data_a = di_a.next()[0]
x_real_a.d = x_data_a
z_style_b.d = z_data_0[j] * (1.0 - r) + z_data_1[j] * r
x_fake_b.forward(clear_buffer=True)
cmp_image = np.concatenate([x_data_a, x_fake_b.d.copy()], axis=3)
images.append(cmp_image)
images = np.concatenate(images)
monitor_image_tile_b.add(i, images)
# Interpolate (B -> A)
z_data_0 = [rng.randn(*z_style_b.shape) for j in range(di_b.size)]
z_data_1 = [rng.randn(*z_style_b.shape) for j in range(di_b.size)]
for i in range(args.num_repeats):
r = 1.0 * i / args.num_repeats
images = []
for j in range(di_b.size):
x_data_b = di_b.next()[0]
x_real_b.d = x_data_b
z_style_a.d = z_data_0[j] * (1.0 - r) + z_data_1[j] * r
x_fake_a.forward(clear_buffer=True)
cmp_image = np.concatenate([x_data_b, x_fake_a.d.copy()], axis=3)
images.append(cmp_image)
images = np.concatenate(images)
monitor_image_tile_a.add(i, images)
def train(args):
# Create Communicator and Context
extension_module = "cudnn"
ctx = get_extension_context(extension_module, type_config=args.type_config)
comm = C.MultiProcessDataParalellCommunicator(ctx)
comm.init()
n_devices = comm.size
mpi_rank = comm.rank
mpi_local_rank = comm.local_rank
device_id = mpi_local_rank
ctx.device_id = str(device_id)
nn.set_default_context(ctx)
# Input
b, c, h, w = args.batch_size, 3, args.image_size, args.image_size
x_real_a = nn.Variable([b, c, h, w])
x_real_b = nn.Variable([b, c, h, w])
# Model
# workaround for starting with the same model among devices.
np.random.seed(412)
maps = args.maps
# within-domain reconstruction (domain A)
x_content_a = content_encoder(x_real_a, maps, name="content-encoder-a")
x_style_a = style_encoder(x_real_a, maps, name="style-encoder-a")
x_recon_a = decoder(x_content_a, x_style_a, name="decoder-a")
# within-domain reconstruction (domain B)
x_content_b = content_encoder(x_real_b, maps, name="content-encoder-b")
x_style_b = style_encoder(x_real_b, maps, name="style-encoder-b")
x_recon_b = decoder(x_content_b, x_style_b, name="decoder-b")
# generate over domains and reconstruction of content and style (domain A)
z_style_a = F.randn(shape=x_style_a.shape)
x_fake_a = decoder(x_content_b, z_style_a, name="decoder-a")
x_content_rec_b = content_encoder(x_fake_a, maps, name="content-encoder-a")
x_style_rec_a = style_encoder(x_fake_a, maps, name="style-encoder-a")
# generate over domains and reconstruction of content and style (domain B)
z_style_b = F.randn(shape=x_style_b.shape)
x_fake_b = decoder(x_content_a, z_style_b, name="decoder-b")
x_content_rec_a = content_encoder(x_fake_b, maps, name="content-encoder-b")
x_style_rec_b = style_encoder(x_fake_b, maps, name="style-encoder-b")
# discriminate (domain A)
p_x_fake_a_list = discriminators(x_fake_a)
p_x_real_a_list = discriminators(x_real_a)
p_x_fake_b_list = discriminators(x_fake_b)
p_x_real_b_list = discriminators(x_real_b)
# Loss
# within-domain reconstruction
loss_recon_x_a = recon_loss(x_recon_a, x_real_a).apply(persistent=True)
loss_recon_x_b = recon_loss(x_recon_b, x_real_b).apply(persistent=True)
# content and style reconstruction
loss_recon_x_style_a = recon_loss(x_style_rec_a,
z_style_a).apply(persistent=True)
loss_recon_x_content_b = recon_loss(x_content_rec_b,
x_content_b).apply(persistent=True)
loss_recon_x_style_b = recon_loss(x_style_rec_b,
z_style_b).apply(persistent=True)
loss_recon_x_content_a = recon_loss(x_content_rec_a,
x_content_a).apply(persistent=True)
# adversarial
def f(x, y):
return x + y
loss_gen_a = reduce(f, [lsgan_loss(p_f)
for p_f in p_x_fake_a_list]).apply(persistent=True)
loss_dis_a = reduce(f, [
lsgan_loss(p_f, p_r)
for p_f, p_r in zip(p_x_fake_a_list, p_x_real_a_list)
]).apply(persistent=True)
loss_gen_b = reduce(f, [lsgan_loss(p_f)
for p_f in p_x_fake_b_list]).apply(persistent=True)
loss_dis_b = reduce(f, [
lsgan_loss(p_f, p_r)
for p_f, p_r in zip(p_x_fake_b_list, p_x_real_b_list)
]).apply(persistent=True)
# loss for generator-related models
loss_gen = loss_gen_a + loss_gen_b \
+ args.lambda_x * (loss_recon_x_a + loss_recon_x_b) \
+ args.lambda_c * (loss_recon_x_content_a + loss_recon_x_content_b) \
+ args.lambda_s * (loss_recon_x_style_a + loss_recon_x_style_b)
# loss for discriminators
loss_dis = loss_dis_a + loss_dis_b
# Solver
lr_g, lr_d, beta1, beta2 = args.lr_g, args.lr_d, args.beta1, args.beta2
# solver for generator-related models
solver_gen = S.Adam(lr_g, beta1, beta2)
with nn.parameter_scope("generator"):
params_gen = nn.get_parameters()
solver_gen.set_parameters(params_gen)
# solver for discriminators
solver_dis = S.Adam(lr_d, beta1, beta2)
with nn.parameter_scope("discriminators"):
params_dis = nn.get_parameters()
solver_dis.set_parameters(params_dis)
# Monitor
monitor = Monitor(args.monitor_path)
# time
monitor_time = MonitorTimeElapsed("Training time", monitor, interval=10)
# reconstruction
monitor_loss_recon_x_a = MonitorSeries("Recon Loss Image A",
monitor,
interval=10)
monitor_loss_recon_x_content_b = MonitorSeries("Recon Loss Content B",
monitor,
interval=10)
monitor_loss_recon_x_style_a = MonitorSeries("Recon Loss Style A",
monitor,
interval=10)
monitor_loss_recon_x_b = MonitorSeries("Recon Loss Image B",
monitor,
interval=10)
monitor_loss_recon_x_content_a = MonitorSeries("Recon Loss Content A",
monitor,
interval=10)
monitor_loss_recon_x_style_b = MonitorSeries("Recon Loss Style B",
monitor,
interval=10)
# adversarial
monitor_loss_gen_a = MonitorSeries("Gen Loss A", monitor, interval=10)
monitor_loss_dis_a = MonitorSeries("Dis Loss A", monitor, interval=10)
monitor_loss_gen_b = MonitorSeries("Gen Loss B", monitor, interval=10)
monitor_loss_dis_b = MonitorSeries("Dis Loss B", monitor, interval=10)
monitor_losses = [
# reconstruction
(monitor_loss_recon_x_a, loss_recon_x_a),
(monitor_loss_recon_x_content_b, loss_recon_x_content_b),
(monitor_loss_recon_x_style_a, loss_recon_x_style_a),
(monitor_loss_recon_x_b, loss_recon_x_b),
(monitor_loss_recon_x_content_a, loss_recon_x_content_a),
(monitor_loss_recon_x_style_b, loss_recon_x_style_b),
# adaversarial
(monitor_loss_gen_a, loss_gen_a),
(monitor_loss_dis_a, loss_dis_a),
(monitor_loss_gen_b, loss_gen_b),
(monitor_loss_dis_b, loss_dis_b)
]
# image
monitor_image_a = MonitorImage("Fake Image B to A Train",
monitor,
interval=1)
monitor_image_b = MonitorImage("Fake Image A to B Train",
monitor,
interval=1)
monitor_images = [
(monitor_image_a, x_fake_a),
(monitor_image_b, x_fake_b),
]
# DataIterator
rng_a = np.random.RandomState(device_id)
rng_b = np.random.RandomState(device_id + n_devices)
di_a = munit_data_iterator(args.img_path_a, args.batch_size, rng=rng_a)
di_b = munit_data_iterator(args.img_path_b, args.batch_size, rng=rng_b)
# Train
for i in range(args.max_iter // n_devices):
ii = i * n_devices
# Train generator-related models
x_data_a, x_data_b = di_a.next()[0], di_b.next()[0]
x_real_a.d, x_real_b.d = x_data_a, x_data_b
solver_gen.zero_grad()
loss_gen.forward(clear_no_need_grad=True)
loss_gen.backward(clear_buffer=True)
comm.all_reduce([w.grad for w in params_gen.values()])
solver_gen.weight_decay(args.weight_decay_rate)
solver_gen.update()
# Train discriminators
x_data_a, x_data_b = di_a.next()[0], di_b.next()[0]
x_real_a.d, x_real_b.d = x_data_a, x_data_b
x_fake_a.need_grad, x_fake_b.need_grad = False, False
solver_dis.zero_grad()
loss_dis.forward(clear_no_need_grad=True)
loss_dis.backward(clear_buffer=True)
comm.all_reduce([w.grad for w in params_dis.values()])
solver_dis.weight_decay(args.weight_decay_rate)
solver_dis.update()
x_fake_a.need_grad, x_fake_b.need_grad = True, True
# LR schedule
if (i + 1) % (args.lr_decay_at_every // n_devices) == 0:
lr_d = solver_dis.learning_rate() * args.lr_decay_rate
lr_g = solver_gen.learning_rate() * args.lr_decay_rate
solver_dis.set_learning_rate(lr_d)
solver_gen.set_learning_rate(lr_g)
if mpi_local_rank == 0:
# Monitor
monitor_time.add(ii)
for mon, loss in monitor_losses:
mon.add(ii, loss.d)
# Save
if (i + 1) % (args.model_save_interval // n_devices) == 0:
for mon, x in monitor_images:
mon.add(ii, x.d)
nn.save_parameters(
os.path.join(args.monitor_path,
"param_{:05d}.h5".format(i)))
if mpi_local_rank == 0:
# Monitor
for mon, loss in monitor_losses:
mon.add(ii, loss.d)
# Save
for mon, x in monitor_images:
mon.add(ii, x.d)
nn.save_parameters(
os.path.join(args.monitor_path, "param_{:05d}.h5".format(i)))