def optimize_latents():
print("Optimizing Latents.")
synthesizer = StyleGANGenerator(args.model_type).model.synthesis
latent_optimizer = LatentOptimizer(synthesizer, args.vgg_layer)
# Optimize only the dlatents.
for param in latent_optimizer.parameters():
param.requires_grad_(False)
if args.video or args.save_optimized_image:
# Hook, saves an image during optimization to be used to create video.
generated_image_hook = GeneratedImageHook(
latent_optimizer.post_synthesis_processing, args.save_frequency)
reference_image = load_images([args.image_path])
reference_image = torch.from_numpy(reference_image).cuda()
reference_image = latent_optimizer.vgg_processing(reference_image)
reference_features = latent_optimizer.vgg16(reference_image).detach()
reference_image = reference_image.detach()
if args.use_latent_finder:
image_to_latent = ImageToLatent().cuda()
image_to_latent.load_state_dict(torch.load(args.image_to_latent_path))
image_to_latent.eval()
latents_to_be_optimized = image_to_latent(reference_image)
latents_to_be_optimized = latents_to_be_optimized.detach().cuda(
).requires_grad_(True)
else:
latents_to_be_optimized = torch.zeros(
(1, 18, 512)).cuda().requires_grad_(True)
criterion = LatentLoss()
optimizer = torch.optim.SGD([latents_to_be_optimized],
lr=args.learning_rate)
progress_bar = tqdm(range(args.iterations))
for step in progress_bar:
optimizer.zero_grad()
generated_image_features = latent_optimizer(latents_to_be_optimized)
loss = criterion(generated_image_features, reference_features)
loss.backward()
loss = loss.item()
optimizer.step()
progress_bar.set_description("Step: {}, Loss: {}".format(step, loss))
optimized_dlatents = latents_to_be_optimized.detach().cpu().numpy()
np.save(args.dlatent_path, optimized_dlatents)
if args.video:
images_to_video(generated_image_hook.get_images(), args.video_path)
if args.save_optimized_image:
save_image(generated_image_hook.last_image, args.optimized_image_path)
def optimize_latents():
# logger
logger = make_logger("project", args.output_dir, 'log')
logger.info("Optimizing Latents.") # optimize to latent
global_time = time.time()
# the module generating image
synthesizer = StyleGANGenerator(args.model_type).model.synthesis
# the optimizer severs as getting the featuremap of the image
latent_optimizer = LatentOptimizer(synthesizer, args.vgg_layer)
# Optimizer only the dlatents.
for param in latent_optimizer.parameters(): # frozen the parameters
param.requires_grad_(False)
if args.video or args.save_optimized_image: # if need save video or the optimized_image
# Hook, saves an image during optimization to be used to create video.
generated_image_hook = GeneratedImageHook(
latent_optimizer.post_synthesis_processing, args.save_frequency)
image_to_latent = ImageToLatent().cuda()
FeatureMapper = FeatureMapping().cuda()
FeatureMapper.load_state_dict(
torch.load('./output/models/FeatureMapper_90.pth'))
#image_to_latent = torch.load('./output/models/ImageToLatent_105.pth')
# best 110
# try below 100
image_to_latent.load_state_dict(
torch.load('output/models/ImageToLatent_90.pth'))
# for param in image_to_latent.parameters(): # frozen the parameters
# param.requires_grad_(False)
# You can use models from scatch
#image_to_latent.train() # training for the first step
FeatureMapper.train() # The feature mapping step
X_train, Y_train = DataSet()
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(image_to_latent.parameters(),
lr=args.learning_rate)
save_dir = os.path.join(args.output_dir, 'models')
os.makedirs(save_dir, exist_ok=True)
device = torch.device('cuda')
imgprocess = ImageProcessing().cuda().eval()
imgpostprocess = PostSynthesisProcessing().cuda().eval()
sum_loss = 0
batch_size = 8
for epoch in range(1, args.epochs + 1):
# logger.info("Epoch: [%d]" % epoch)
index = [i for i in range(len(X_train))]
random.shuffle(index)
X_train = X_train[index]
Y_train = Y_train[index]
index_b = []
ite = int(len(X_train) / batch_size)
for i in range(ite):
batchx = index[i * batch_size:i * batch_size + batch_size]
batchy = Y_train[i * batch_size:i * batch_size + batch_size]
index_b.append(batchx)
for (i, batchind) in enumerate(index_b, 1):
batchimg = X_train[batchind].astype(np.float16)
batchlats = Y_train[batchind].astype(np.float16)
reference_image = torch.from_numpy(batchimg).float().to(device)
reference_image = reference_image.permute(0, 3, 1, 2)
generated_features = torch.from_numpy(batchlats).float().to(device)
latents_optimized = image_to_latent(reference_image)
# On first step , this do not use
latents_to_be_optimized = latents_optimized.view(-1, 18 * 512)
latents_to_be_optimized = FeatureMapper(latents_to_be_optimized)
latents_to_be_optimized = latents_to_be_optimized.view(-1, 18, 512)
latents_to_be_optimized = latents_to_be_optimized.cuda(
).requires_grad_(True)
generated_imgs = latent_optimizer(latents_to_be_optimized,
batchimg)
nimgs = latent_optimizer(generated_features, batchimg)
gen_feat = latent_optimizer.vgg16(imgprocess(generated_imgs))
n_feat = latent_optimizer.vgg16(imgprocess(nimgs))
# until here , for the second step (Mapping)
# loss=criterion(latents_to_be_optimized,generated_features)
loss = criterion(gen_feat, n_feat) + criterion(
generated_imgs, nimgs) # second loss for second step
optimizer.zero_grad()
loss.backward()
optimizer.step()
sum_loss += loss.item()
if (epoch) % 1 == 0:
print("epoch [" + str(epoch) + "] finished!")
if (epoch) % 10 == 0:
elapsed = time.time() - global_time
elapsed = str(datetime.timedelta(seconds=elapsed)).split('.')[0]
logger.info("Elapsed: [%s] Epoch: [%d] Step: %d Loss: %f" %
(elapsed, epoch, i, sum_loss / 300))
sum_loss = 0
model_save_file = os.path.join(
save_dir, "FeatureMapper" + "_" + str(epoch) + ".pth")
torch.save(FeatureMapper.state_dict(), model_save_file)
model_save_file = os.path.join(
save_dir, "ImageToLatent" + "_" + str(epoch) + ".pth")
torch.save(image_to_latent.state_dict(), model_save_file)
image_dir = args.optimized_image_path + \
str(epoch) + "_" + ".jpg"
save_img = generated_imgs[0].detach().cpu().numpy()
save_image(save_img, image_dir)
image_dir = args.optimized_image_path + \
str(epoch) + "_r" + ".jpg"
save_image(nimgs[0].detach().cpu().numpy(), image_dir)
if args.video:
images_to_video(generated_image_hook.get_images(), args.video_path)
if args.save_optimized_image:
save_image(generated_image_hook.last_image, args.optimized_image_path)
def optimize_latents():
# logger
logger = make_logger("project", args.output_dir, 'log')
logger.info("Optimizing Latents.") # optimize to latent
global_time = time.time()
# the module generating image
synthesizer = StyleGANGenerator(args.model_type).model.synthesis
# the optimizer severs as getting the featuremap of the image
latent_optimizer = LatentOptimizer(synthesizer, args.vgg_layer)
# Optimizer only the dlatents.
for param in latent_optimizer.parameters(): # frozen the parameters
param.requires_grad_(False)
if args.video or args.save_optimized_image: # if need save video or the optimized_image
# Hook, saves an image during optimization to be used to create video.
generated_image_hook = GeneratedImageHook(
latent_optimizer.post_synthesis_processing, args.save_frequency)
image_to_latent = ImageToLatent().cuda()
# image_to_latent.load_state_dict(torch.load('output/models'))
image_to_latent.train()
dataset = make_dataset(folder=False,
img_dir='./data/sy_imgs',
resolution=1024)
data = get_data_loader(dataset, batch_size=16, num_workers=4)
criterion = LatentLoss()
optimizer = torch.optim.Adam(image_to_latent.parameters(),
lr=args.learning_rate)
save_dir = os.path.join(args.output_dir, 'models')
os.makedirs(save_dir, exist_ok=True)
device = torch.device('cuda')
imgprocess = ImageProcessing().cuda().eval()
imgpostprocess = PostSynthesisProcessing().cuda().eval()
sum_loss = 0
for epoch in range(1, args.epochs + 1):
# logger.info("Epoch: [%d]" % epoch)
for (i, batch) in enumerate(data, 1):
reference_image = batch.to(device)
style_image = reference_image.clone() # [1, 3, 1024, 2024]
reference_image = imgprocess(reference_image)
# print(reference_image.size())
# print(reference_image[0][1]) # 值在0到1之间 tensor([-0.2157, -0.2157, -0.2235, ..., 0.0431, 0.0431, 0.0431]
# reference_image = latent_optimizer.vgg_processing(
# reference_image) # normalization
# print(reference_image.size()) # ([1, 3, 256, 256])
# print(reference_image[0][1][1]) # 这里出了问题，去了和上面一样的tensor，基本都变成了-2.03
# assert False
reference_features = latent_optimizer.vgg16(reference_image)
# reference_image = reference_image.detach()
# latents should be get from the gmapping class from stylegan
latents_to_be_optimized = image_to_latent(reference_image)
# print(latents_to_be_optimized.size()) # torch.Size([1, 18, 512])
# print(latents_to_be_optimized)
# latents_to_be_optimized = latents_to_be_optimized.detach().cuda().requires_grad_(True)
latents_to_be_optimized = latents_to_be_optimized.cuda(
).requires_grad_(True)
generated_image = latent_optimizer(latents_to_be_optimized,
style_image)
generated_image_features = latent_optimizer.vgg16(generated_image)
# print(generated_image_features.size()) # torch.Size([1, 3, 256, 256])
# print(generated_image_features[0][1])
# assert False
loss = criterion(generated_image_features, reference_features)
optimizer.zero_grad()
loss.backward()
optimizer.step()
sum_loss += loss.item()
if (epoch) % 1 == 0:
elapsed = time.time() - global_time
elapsed = str(datetime.timedelta(seconds=elapsed)).split('.')[0]
logger.info("Elapsed: [%s] Epoch: [%d] Step: %d Loss: %f" %
(elapsed, epoch, i, sum_loss / 100))
print(latents_to_be_optimized[0, :10, :10])
print()
sum_loss = 0
model_save_file = os.path.join(
save_dir, "ImageToLatent" + "_" + str(epoch) + ".pth")
torch.save(image_to_latent.state_dict(), model_save_file)
image_dir = args.optimized_image_path + \
str(epoch) + "_" + ".jpg"
save_img = imgpostprocess(
generated_image[0].detach().cpu()).numpy()
save_image(save_img, image_dir)
if args.video:
images_to_video(generated_image_hook.get_images(), args.video_path)
if args.save_optimized_image:
save_image(generated_image_hook.last_image, args.optimized_image_path)
def optimize_latents():
print("Optimizing Latents.")
generator = StyledGenerator(512).to(device)
generator.load_state_dict(torch.load(args.path)['generator'])
generator.eval()
latent_optimizer = LatentOptimizer(generator, args.vgg_layer)
mean_style = get_mean_style(generator, device)
total = np.zeros((83 * 3, 512))
# Optimize only the dlatents.
for param in latent_optimizer.parameters():
param.requires_grad_(False)
if args.video or args.save_optimized_image:
# Hook, saves an image during optimization to be used to create video.
generated_image_hook = GeneratedImageHook(
latent_optimizer.post_synthesis_processing, args.save_frequency)
for i in range(3 * 83): #3 for each pictrue
iid = i % 3
path = int(i / 3)
iterations = int(200 * iid + 300)
image_path = './data/' + str(path) + '.jpg'
print(image_path)
reference_image = load_images([image_path])
reference_image = torch.from_numpy(reference_image).to(device)
reference_image = latent_optimizer.vgg_processing(
reference_image) #normalize
reference_features = latent_optimizer.vgg16(
reference_image).detach() #vgg
reference_image = reference_image.detach()
if args.use_latent_finder:
image_to_latent = ImageToLatent().cuda()
image_to_latent.load_state_dict(
torch.load(args.image_to_latent_path))
image_to_latent.eval()
latents_to_be_optimized = image_to_latent(reference_image)
latents_to_be_optimized = latents_to_be_optimized.detach().cuda(
).requires_grad_(True)
else:
latents_to_be_optimized = torch.zeros(
(1, 512)).cuda().requires_grad_(True)
criterion = LatentLoss()
optimizer = torch.optim.SGD([latents_to_be_optimized],
lr=args.learning_rate)
progress_bar = tqdm(range(iterations))
for step in progress_bar:
optimizer.zero_grad()
generated_image_features = latent_optimizer(
latents_to_be_optimized, mean_style, i)
#print(latents_to_be_optimized)
loss = criterion(generated_image_features, reference_features)
loss.backward()
loss = loss.item()
optimizer.step()
progress_bar.set_description("Step: {}, Loss: {}".format(
step, loss))
optimized_dlatents = latents_to_be_optimized.detach().cpu().numpy()
total[i] = optimized_dlatents[0]
np.save(args.dlatent_path, total)
def optimize_latents():
print("Optimizing Latents.")
generator = StyledGenerator(512).to(device)
generator.load_state_dict(torch.load(args.path)['generator'])
generator.eval()
latent_optimizer = LatentOptimizer(generator, args.vgg_layer)
mean_style = get_mean_style(generator, device)
# Optimize only the dlatents.
for param in latent_optimizer.parameters():
param.requires_grad_(False)
if args.video or args.save_optimized_image:
# Hook, saves an image during optimization to be used to create video.
generated_image_hook = GeneratedImageHook(
latent_optimizer.post_synthesis_processing, args.save_frequency)
reference_image = load_images([args.image_path])
reference_image = torch.from_numpy(reference_image).to(device)
reference_image = latent_optimizer.vgg_processing(
reference_image) #normalize
utils.save_image(reference_image,
'./reference.png',
nrow=1,
normalize=True,
range=(-1, 1))
reference_features = latent_optimizer.vgg16(reference_image).detach() #vgg
reference_image = reference_image.detach()
if args.use_latent_finder:
image_to_latent = ImageToLatent().cuda()
image_to_latent.load_state_dict(torch.load(args.image_to_latent_path))
image_to_latent.eval()
latents_to_be_optimized = image_to_latent(reference_image)
latents_to_be_optimized = latents_to_be_optimized.detach().cuda(
).requires_grad_(True)
else:
latents_to_be_optimized = torch.zeros(
(1, 512)).cuda().requires_grad_(True)
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
criterion = LatentLoss()
optimizer = torch.optim.RMSprop([latents_to_be_optimized],
lr=args.learning_rate,
weight_decay=0.02)
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
progress_bar = tqdm(range(args.iterations))
for step in progress_bar:
#print(latents_to_be_optimized)
optimizer.zero_grad()
generated_image_features = latent_optimizer(latents_to_be_optimized,
mean_style)
loss = criterion(generated_image_features, reference_features)
loss.backward()
loss = loss.item()
optimizer.step()
# if step==args.iterations:
# break
# with torch.no_grad():
# latents_to_be_optimized.add_(-latents_to_be_optimized.mean()+3e-2*torch.randn(1).to('cuda'))
# latents_to_be_optimized.div_(latents_to_be_optimized.std()+3e-2*torch.randn(1).to('cuda'))
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
progress_bar.set_description("Step: {}, Loss: {}".format(step, loss))
print(latents_to_be_optimized)
#latents_to_be_optimized=latent_optimizer.normalize(latents_to_be_optimized)
#print(latents_to_be_optimized.mean(),latents_to_be_optimized.std())
optimized_dlatents = latents_to_be_optimized.detach().cpu().numpy()
np.save(args.dlatent_path, optimized_dlatents)
if args.video:
images_to_video(generated_image_hook.get_images(), args.video_path)
if args.save_optimized_image:
save_image(generated_image_hook.last_image, args.optimized_image_path)