命令
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_transform self.target_transform = target_transform self.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_im
return 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_avgshape 是 （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_avgshape（16，512）images, result_latent = self.decoder([codes], input_is_latent=input_is_latent, randomize_noise=randomize_noise, return_latents=return_latents)解码器输出预测图片。
其中noiselist：15，out是图像特征 shape（4，64，512，512）decoder本身的参数如下
返回的imageshape(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()
