Lama:《Resolution-robust Large Mask Inpainting with Fourier Convolutions》训练、推理实战记录

记录一下Lama模型的训练、infe踩坑,以及如何更改预设的mask生成方式。
一、环境
简单提一下,一定要按照作者给的requirements.txt里的库版本安装,hydra-core和pytorch-lightning最新版本在此项目代码上均会报错无法运行。

二、预训练模型微调
lama的训练全部是以配置文件.yaml的方式进行的,所以针对不同数据集的预训练模型所使用的yaml也是不同的。总体上作者是在PLACES和CelebA上进行了预训练,同时也包含了Big-Lama、Lama-fourier等多种模型细节的配置。就以在Places-Challenge效果最好的Big-Lama为例,在预训练权重上使用自己的数据和mask进行微调。
1、定义自己的mask生成方式
假如训练中不想使用原本预设的mask生成方式,想要根据自己的任务进行生成(例如在特定位置以特定形状进行生成mask)。
在saicinpainting/training/data/masks.py中首先参照作者方式构建自己的mask生成类。例如RandomSuperresMaskGenerator类是执行make_random_superres_mask函数返回一个矩阵mask,代表覆盖区域。
Lama:《Resolution-robust Large Mask Inpainting with Fourier Convolutions》训练、推理实战记录_第1张图片

之后在同文件中的get_mask_generator函数中增添自己的类对应的kind。
Lama:《Resolution-robust Large Mask Inpainting with Fourier Convolutions》训练、推理实战记录_第2张图片

最后在saicinpainting/training/data/datasets.py的make_default_train_dataloader函数中通过指定get_mask_generator对应的kind参数实现自定义mask类的调用。这个在训练过程中可以在config中指定
2、数据准备
在官方仓库的readme中,根据Train and Eval下的create your data顺序进行即可。其中不一定非要使用bin/gen_mask_dataset.py来进行数据生成,完全可以使用任何自定义的脚本,只要最后形成相同的文件目录和对应文件即可。
3、train
原本作者的big-lama预训练权重中缺少了discriminitor的参数,且有部分参数冗余,所以直接一步到位按如下链接下载https://disk.yandex.ru/d/wJ2Ee0f1HvasDQ/big-lama-with-discr。
之后只需要更改config.yaml对应的各种路径即可实现训练。

你可能感兴趣的:(python,深度学习,机器学习,生成对抗网络,计算机视觉)