【源码解析】psp training 过程

命令

python scripts/train.py   \
--dataset_type=ffhq_encode \
--exp_dir=./logs/mytraining \
--workers=8 \
--batch_size=8 \
--test_batch_size=8 \
--test_workers=8 \
--val_interval=2500 \
--save_interval=5000 \
--encoder_type=BackboneEncoderUsingLastLayerIntoWPlus\
--start_from_latent_avg \
--lpips_lambda=0.8 \
--l2_lambda=1 \
--id_lambda=0.1

整合为单行为python scripts/train.py --dataset_type=ffhq_encode --exp_dir=./logs/mytraining --workers=8 --batch_size=8 --test_batch_size=8 --test_workers=8 --val_interval=2500 --save_interval=5000 --encoder_type=BackboneEncoderUsingLastLayerIntoWPlus --start_from_latent_avg --lpips_lambda=0.8 --l2_lambda=1 --id_lambda=0.1

尝试Downloading: "https://raw.githubusercontent.com/richzhang/PerceptualSimilarity/master/lpips/weights/v0.1/alex.pth" to /home/joselyn/.cache/torch/hub/checkpoints/alex.pth 然后报错：urllib.error.URLError: <urlopen error [Errno 111] Connection refused>

更换encoder_type 再执行python scripts/train.py --dataset_type=ffhq_encode --exp_dir=./logs/mytraining --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
仍然报相同的错误。而且还找不到这个是在哪里执行的。

找到了。在coach 初始化中

     if self.opts.lpips_lambda > 0:self.lpips_loss = LPIPS(net_type='alex').to(self.device).eval()

初始化lpips损失的时候，需要下载。

再次尝试运行： --dataset_type=ffhq_encode --exp_dir=./logs/mytraining --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

报错RuntimeError: CUDA out of memory. Tried to allocate 1.00 GiB (GPU 0; 10.76 GiB total capacity; 8.43 GiB already allocated; 934.44 MiB free; 8.71 GiB reserved in total by PyTorch)
本机是2080ti 12G，接下来调整参数如下：--dataset_type=ffhq_encode --exp_dir=../logs/mytraining --workers=2 --batch_size=4 --test_batch_size=4 --test_workers=2 --val_interval=2500 --save_interval=5000 --encoder_type=GradualStyleEncoder --start_from_latent_avg --lpips_lambda=0.8 --l2_lambda=1 --id_lambda=0.1，再次尝试，仍然内存不足。

train-debug过程

train_dataset = ImagesDataset(source_root=dataset_args['train_source_root'], target_root=dataset_args['train_target_root'], source_transform=transforms_dict['transform_source'], target_transform=transforms_dict['transform_gt_train'], opts=self.opts)
- ImagesDataset的构造函数为
```
 def __init__(self, source_root, target_root, opts, target_transform=None, source_transform=None):self.source_paths = sorted(data_utils.make_dataset(source_root))self.target_paths = sorted(data_utils.make_dataset(target_root))self.source_transform = source_transformself.target_transform = target_transformself.opts = opts
```
- 计算结果如下：
- 对应的测试集如下：
coach.train() 开始执行时，coach的参数信息如下：
下面进入整体的训练循环中。

     while self.global_step < self.opts.max_steps:for batch_idx, batch in enumerate(self.train_dataloader):

batch_idx = 0, batch 是2张图片的list。为什么是2，batch_size不是4吗？

这里来回答一下为什么是2。
两个分表代表训练集的sample和label。这个是由imageDataset 类中__getitem__函数决定的。具体如下：

 def __getitem__(self, index):from_path = self.source_paths[index]from_im = Image.open(from_path)from_im = from_im.convert('RGB') if self.opts.label_nc == 0 else from_im.convert('L')to_path = self.target_paths[index]to_im = Image.open(to_path).convert('RGB')if self.target_transform:to_im = self.target_transform(to_im)if self.source_transform:from_im = self.source_transform(from_im)else:from_im = to_imreturn from_im, to_im  # 返回的是x,y

每个元素是这样的，shape是（4，3，256，256），其中4是batch_size的值。

获取x和y： x, y = batch ， x, y = x.to(self.device).float(), y.to(self.device).float()
把x给psp 网络，进行前向计算： y_hat, latent = self.net.forward(x, return_latents=True)

下面我们进入到psp的前向计算过程

psp 网络的参数如下：
其中latent_avg shape 是（16，512）为什么这里会是16呢？
codes = self.encoder(x) 获取编码器的结果。codes（tensor：[batch_size, 16,512]）
if start_from_latent_avg and not learn_in_w ： codes = codes + self.latent_avg.repeat(codes.shape[0], 1, 1) 。 codes（tensor：[batch_size, 16,512]），self.latent_avg shape（16，512）
images, result_latent = self.decoder([codes], input_is_latent=input_is_latent, randomize_noise=randomize_noise, return_latents=return_latents) 解码器输出预测图片。
其中noise list：15，out 是图像特征 shape（4，64，512，512） decoder本身的参数如下
返回的image shape(batch_size,3,512,512), 返回的latent 还是（4，16，512），但latent不是原来的latent了，他经过了如下的计算：

latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1)latent = torch.cat([latent, latent2], 1)

回到psp的前向过程

decoder得到的image在resize的情况下，进行images = self.face_pool(images) 。得到image 的shape为[4,3,256,256]

返回到coach 的train函数中

y_hat, latent = self.net.forward(x, return_latents=True) 得到了预测的图像和最终的latent_code。
loss, loss_dict, id_logs = self.calc_loss(x, y, y_hat, latent) 计算损失。
后面的就可以不用管了。

CelebAMask-HQ 下载地址

https://github.com/switchablenorms/CelebAMask-HQ

划分数据集。参考https://www.codeleading.com/article/6296310738/

transforms_dict = dataset_args['transforms'](self.opts).get_transforms()