【源码解析】Encoding in Style: a StyleGAN Encoder for Image-to-Image Translation
文章目录
- Readme
-
- training：
- inference
- inference.py
-
- 参数
- 人脸对齐工作
- 尝试进行inference
- 可视化潜码
- 问题
Readme
- 发现了http://localhost:8888/notebooks/notebooks/inference_playground.ipynb 待用
- 预训练模型下载中，要使用预训练模型
--checkpoint_path
training：
- 如果要在自己的数据集上面实验， 在
data_configs.py总定义数据路径。 在transforms_configs.py中定义数据transforms。
例子：
dataset_paths = {
'ffhq': '/path/to/ffhq/images256x256' # 这个是数据集还是文件
'celeba_test': '/path/to/CelebAMask-HQ/test_img', # 这是测试的图片？
}
- 主要的训练文件是
scripts/train.py.。 - 中间训练结果保存在
opts.exp_dir下面，包含了 checkpoint，训练输出和测试输出。
训练psp encoder的命令
python scripts/train.py \
--dataset_type=ffhq_encode \
--exp_dir=/path/to/experiment \
--workers=8 \
--batch_size=8 \
--test_batch_size=8 \
--test_workers=8 \
--val_interval=2500 \
--save_interval=5000 \
--encoder_type=GradualStyleEncoder \
--start_from_latent_avg \
--lpips_lambda=0.8 \
--l2_lambda=1 \
--id_lambda=0.1
- By default, we assume that the StyleGAN used outputs images at resolution 1024x1024. If you wish to use a StyleGAN at a smaller resolution, you can do so by using the flag
--output_size (e.g., --output_size=256).
inference
scripts/inference.py- example
python scripts/inference.py \
--exp_dir=/path/to/experiment \
--checkpoint_path=experiment/checkpoints/best_model.pt \
--data_path=/path/to/test_data \
--test_batch_size=4 \
--test_workers=4 \
--couple_outputs
给数据集还是给单张图片？
- 调整输入图像的分辨率。
--resize_factors - 调整输出图像的分辨率。
--resize_outputs
inference.py
参数
# arguments for inference script
self.parser.add_argument('--exp_dir', type=str, help='Path to experiment output directory') # 指定一个实验特定的输出路径。
self.parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to pSp model checkpoint') # 用哪个模型进行推断
self.parser.add_argument('--data_path', type=str, default='gt_images', help='Path to directory of images to evaluate')
self.parser.add_argument('--couple_outputs', action='store_true', help='Whether to also save inputs + outputs side-by-side')
self.parser.add_argument('--resize_outputs', action='store_true', help='Whether to resize outputs to 256x256 or keep at 1024x1024') # 加上这个参数，我输出分辨率就是256
self.parser.add_argument('--test_batch_size', default=2, type=int, help='Batch size for testing and inference')
self.parser.add_argument('--test_workers', default=2, type=int, help='Number of test/inference dataloader workers')
# arguments for style-mixing script
self.parser.add_argument('--n_images', type=int, default=None, help='Number of images to output. If None, run on all data') # 输出图像的数量。
self.parser.add_argument('--n_outputs_to_generate', type=int, default=5, help='Number of outputs to generate per input image.') # 每个输入图像生出的输出数
self.parser.add_argument('--mix_alpha', type=float, default=None, help='Alpha value for style-mixing')
self.parser.add_argument('--latent_mask', type=str, default=None, help='Comma-separated list of latents to perform style-mixing with') # 要执行样式混合的用逗号分隔的隐层列表
# arguments for super-resolution
self.parser.add_argument('--resize_factors', type=str, default=None,
help='Downsampling factor for super-res (should be a single value for inference).')
为了方便理解，再看一下train_options里面的参数。
self.parser.add_argument('--exp_dir', type=str, help='Path to experiment output directory')
self.parser.add_argument('--dataset_type', default='ffhq_encode', type=str, help='Type of dataset/experiment to run')
self.parser.add_argument('--encoder_type', default='GradualStyleEncoder', type=str, help='Which encoder to use') # 编码类型
self.parser.add_argument('--input_nc', default=3, type=int, help='Number of input image channels to the psp encoder') # 编码器对应的输入图像通道。
self.parser.add_argument('--label_nc', default=0, type=int, help='Number of input label channels to the psp encoder') # 输入标签通道。 标签是什么？
self.parser.add_argument('--output_size', default=1024, type=int, help='Output size of generator') # 生成器的输出规模
self.parser.add_argument('--batch_size', default=4, type=int, help='Batch size for training')
self.parser.add_argument('--test_batch_size', default=2, type=int, help='Batch size for testing and inference')
self.parser.add_argument('--workers', default=4, type=int, help='Number of train dataloader workers')
self.parser.add_argument('--test_workers', default=2, type=int, help='Number of test/inference dataloader workers')
self.parser.add_argument('--learning_rate', default=0.0001, type=float, help='Optimizer learning rate')
self.parser.add_argument('--optim_name', default='ranger', type=str, help='Which optimizer to use') # 优化起选择
self.parser.add_argument('--train_decoder', default=False, type=bool, help='Whether to train the decoder model') # 是否训练解码器
self.parser.add_argument('--start_from_latent_avg', action='store_true', help='Whether to add average latent vector to generate codes from encoder.')
# 是否加入平均潜在向量从编码器生成代码。
self.parser.add_argument('--learn_in_w', action='store_true', help='Whether to learn in w space instead of w+') # 选择学习空间
# 关于损失的设定
self.parser.add_argument('--lpips_lambda', default=0.8, type=float, help='LPIPS loss multiplier factor')
self.parser.add_argument('--id_lambda', default=0, type=float, help='ID loss multiplier factor')
self.parser.add_argument('--l2_lambda', default=1.0, type=float, help='L2 loss multiplier factor')
self.parser.add_argument('--w_norm_lambda', default=0, type=float, help='W-norm loss multiplier factor')
self.parser.add_argument('--lpips_lambda_crop', default=0, type=float, help='LPIPS loss multiplier factor for inner image region')
self.parser.add_argument('--l2_lambda_crop', default=0, type=float, help='L2 loss multiplier factor for inner image region')
self.parser.add_argument('--moco_lambda', default=0, type=float, help='Moco-based feature similarity loss multiplier factor')
self.parser.add_argument('--stylegan_weights', default=model_paths['stylegan_ffhq'], type=str, help='Path to StyleGAN model weights') # stylegan网络的权重路劲
self.parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to pSp model checkpoint') # 模型路径
self.parser.add_argument('--max_steps', default=500000, type=int, help='Maximum number of training steps') # 最多训练多少步
self.parser.add_argument('--image_interval', default=100, type=int, help='Interval for logging train images during training') # 记录图像的间隔
self.parser.add_argument('--board_interval', default=50, type=int, help='Interval for logging metrics to tensorboard')
self.parser.add_argument('--val_interval', default=1000, type=int, help='Validation interval') # 验证间隔
self.parser.add_argument('--save_interval', default=None, type=int, help='Model checkpoint interval') # 保存checkpoint的间隔。
# arguments for weights & biases support
self.parser.add_argument('--use_wandb', action="store_true", help='Whether to use Weights & Biases to track experiment.')
# 是否使用权重和偏差来跟踪实验
# arguments for super-resolution
self.parser.add_argument('--resize_factors', type=str, default=None, help='For super-res, comma-separated resize factors to use for inference.')
模型下载完成，尝试根据网络图片进行推断。
人脸对齐工作
root = './ooframes/'
imgs = os.listdir(root)
paths = [root+p for p in imgs]
# paths =['./ooframes/test_img.png']
print(len(paths))
with Pool(8) as p:
p.map(align_face, paths)
修改后的图像会保存在aframes目录下。
尝试进行inference
编辑运行命令
python scripts/inference.py \
--exp_dir=./logs/infer_test \
--checkpoint_path=./experiment/checkpoints/psp_ffhq_encode.pt \
--data_path=./aframes \
--test_batch_size=1 \
--test_workers=1 \
--couple_outputs
尝试运行命令， Debug调试
ckpt = torch.load(test_opts.checkpoint_path, map_location='cpu')
- net = pSp(opts)根据参数进行模型的初始化。但还没有用到已经训练过的参数。
dataset_args = data_configs.DATASETS[opts.dataset_type]
transforms_dict = dataset_args['transforms'](opts).get_transforms()不晓得这一步是在做什么？ 在为下一步做铺垫。
dataset = InferenceDataset(root=opts.data_path, transform=transforms_dict['transform_inference'], opts=opts)
for input_batch in tqdm(dataloader)运行报错：
解决方法： 不适用tqdm。
成功输出风格转换后的图像。
可视化潜码
记录一下opts的参数。
- 修改代码，增加潜码的输出。
- net 获得的codes 和images的shape分别为
尝试可视化 code，Image.fromarray(np.array(target_w)).convert("L").save(codes_save_path)
漆黑一片，这里就不写了。
问题
-
ffhq_encode、ffhq_frontaliz、celebs_sketch_to_face 等数据类型有什么区别。
-
net.eval()net.eval() 不启用 BatchNormalization 和 Dropout。此时pytorch会自动把BN和DropOut固定住，不会取平均，而是用训练好的值。不然的话，一旦test的batch_size过小，很容易就会因BN层导致模型performance损失较大；\
model.train() ：启用 BatchNormalization 和 Dropout。 在模型测试阶段使用model.train() 让model变成训练模式，此时 dropout和batch normalization的操作在训练q起到防止网络过拟合的问题。
因此，在使用PyTorch进行训练和测试时一定要记得把实例化的model指定train/eval。
model.eval() 负责改变batchnorm、dropout的工作方式，如在eval()模式下，dropout是不工作的。