JoJoGAN 实践
JoJoGAN: One Shot Face Stylization. 只用一张人脸图片,就能学习其风格,然后迁移到其他图片。训练时长只用 1~2 min 即可。
效果:
主流程:
本文分享了个人在本地环境(非 colab)实践 JoJoGAN 的整个过程。你也可以依照本文上手训练自己喜欢的风格。
准备环境
安装:
conda create -n torch python=3.9 -y
conda activate torch
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch -y
检查:
$ python - <<EOF
import torch, torchvision
print(torch.__version__, torch.cuda.is_available())
EOF
1.10.1 True
准备代码
git clone https://github.com/mchong6/JoJoGAN.git
cd JoJoGAN
pip install tqdm gdown matplotlib scipy opencv-python dlib lpips wandb
# Ninja is required to load C++ extensions
wget https://github.com/ninja-build/ninja/releases/download/v1.10.2/ninja-linux.zip
sudo unzip ninja-linux.zip -d /usr/local/bin/
sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force
然后,将本文提供的几个 *.py
放进 JoJoGAN
目录,从这里获取: https://github.com/ikuokuo/start-deep-learning/tree/master/practice/JoJoGAN 。
download_models.py
: 获取模型generate_faces.py
: 生成人脸stylize.py
: 风格化train.py
: 训练
之后,于训练流程一节,会结合代码,讲述下 JoJoGAN 的工作流程。其他些 *.py
只提下用法,实现就不多说了。
获取模型
python download_models.py
获取模型,如下:
models/
├── arcane_caitlyn_preserve_color.pt
├── arcane_caitlyn.pt
├── arcane_jinx_preserve_color.pt
├── arcane_jinx.pt
├── arcane_multi_preserve_color.pt
├── arcane_multi.pt
├── art.pt
├── disney_preserve_color.pt
├── disney.pt
├── dlibshape_predictor_68_face_landmarks.dat
├── e4e_ffhq_encode.pt
├── jojo_preserve_color.pt
├── jojo.pt
├── jojo_yasuho_preserve_color.pt
├── jojo_yasuho.pt
├── restyle_psp_ffhq_encode.pt
├── stylegan2-ffhq-config-f.pt
├── supergirl_preserve_color.pt
└── supergirl.pt
生成人脸
用 StyleGAN2 预训练模型随机生成人脸,用于测试:
python generate_faces.py -n 5 -s 2000 -o input
使用预训练风格
JoJoGAN 给了 8 个预训练模型,可以一并体验,与文首的效果图一样:
# 预览 JoJoGAN 所有预训练模型 风格化某图片(test_input/iu.jpeg)的效果
python stylize.py -i test_input/iu.jpeg -s all --save-all --show-all
# 使用 JoJoGAN 所有预训练模型 风格化所有生成的测试人脸(input/*)
find ./input -type f -print0 | xargs -0 -i python stylize.py -i {} -s all --save-all
训练自己的风格
首先,准备一张风格图:
之后,开始训练:
python train.py -n yinshi -i style_images/yinshi.jpeg --alpha 1.0 --num_iter 500 --latent_dim 512 --use_wandb --log_interval 50
--use_wandb
时,可查看训练日志:
最后,测试效果:
python stylize.py -i input/girl.jpeg --save-all --show-all --test_style yinshi --test_ckpt output/yinshi.pt --test_ref output/yinshi/style_images_aligned/yinshi.png
训练工作流程
准备风格图片,转为训练数据
将风格图片里的人脸裁减对齐:
# dlib 预测人脸特征点,再裁减对齐
from util import align_face
style_aligned = align_face(img_path)
将风格图片 GAN Inversion 逆映射回预训练模型的隐向量空间(Latent Space):
name, _ = os.path.splitext(os.path.basename(img_path))
style_code_path = os.path.join(latent_dir, f'{name}.pt')
# e4e FFHQ encoder (pSp) > GAN inversion,得到 latent
from e4e_projection import projection
latent = projection(style_aligned, style_code_path, device)
载入 StyleGAN2 模型,训练微调
载入预训练模型:
latent_dim = 512
# 加载预训练模型
original_generator = Generator(1024, latent_dim, 8, 2).to(device)
ckpt = torch.load("models/stylegan2-ffhq-config-f.pt", map_location=lambda storage, loc: storage)
original_generator.load_state_dict(ckpt["g_ema"], strict=False)
# 准备微调的模型
generator = deepcopy(original_generator)
训练可调参数:
# 控制风格强度 [0, 1]
alpha = 1.0
alpha = 1-alpha
# 是否保留原图像色彩
preserve_color = True
# 训练迭代次数(最好 500,Adam 学习率是基于 500 次迭代调优的)
num_iter = 500
# 风格图片 targets 及 latents
targets = ..
latents = ..
进行训练,拟合隐空间。最后保存:
# 准备 LPIPS 计算 loss
lpips_fn = lpips.LPIPS(net='vgg').to(device)
# 准备优化器
g_optim = torch.optim.Adam(generator.parameters(), lr=2e-3, betas=(0, 0.99))
# 哪些层用于交换,用于生成风格化图片
if preserve_color:
id_swap = [7,9,11,15,16,17]
else:
id_swap = list(range(7, generator.n_latent))
# 训练迭代
for idx in tqdm(range(num_iter)):
# 交换层混合风格,并加噪声
mean_w = generator.get_latent(torch.randn([latents.size(0), latent_dim])
.to(device)).unsqueeze(1).repeat(1, generator.n_latent, 1)
in_latent = latents.clone()
in_latent[:, id_swap] = alpha*latents[:, id_swap] + (1-alpha)*mean_w[:, id_swap]
# 以 latent 风格化图片,与目标风格对比
img = generator(in_latent, input_is_latent=True)
loss = lpips_fn(F.interpolate(img, size=(256,256), mode='area'),
F.interpolate(targets, size=(256,256), mode='area')).mean()
# 优化
g_optim.zero_grad()
loss.backward()
g_optim.step()
# 保存权重,完成
torch.save({"g": generator.state_dict()}, save_path)
结语
JoJoGAN 实践下来效果不错。使用本文给到的代码,更容易上手训练自己喜欢的风格,值得试试。
JoJoGAN 实践的更多相关文章
- webp图片实践之路
最近,我们在项目中实践了webp图片,并且抽离出了工具模块,整合到了项目的基础模板中.传闻IOS10也将要支持webp,那么使用webp带来的性能提升将更加明显.估计在不久的将来,webp会成为标配. ...
- Hangfire项目实践分享
Hangfire项目实践分享 目录 Hangfire项目实践分享 目录 什么是Hangfire Hangfire基础 基于队列的任务处理(Fire-and-forget jobs) 延迟任务执行(De ...
- TDD在Unity3D游戏项目开发中的实践
0x00 前言 关于TDD测试驱动开发的文章已经有很多了,但是在游戏开发尤其是使用Unity3D开发游戏时,却听不到特别多关于TDD的声音.那么本文就来简单聊一聊TDD如何在U3D项目中使用以及如何使 ...
- Logstash实践: 分布式系统的日志监控
文/赵杰 2015.11.04 1. 前言 服务端日志你有多重视? 我们没有日志 有日志,但基本不去控制需要输出的内容 经常微调日志,只输出我们想看和有用的 经常监控日志,一方面帮助日志微调,一方面及 ...
- 【大型网站技术实践】初级篇:借助Nginx搭建反向代理服务器
一.反向代理:Web服务器的“经纪人” 1.1 反向代理初印象 反向代理(Reverse Proxy)方式是指以代理服务器来接受internet上的连接请求,然后将请求转发给内部网络上的服务器,并将从 ...
- Windows平台分布式架构实践 - 负载均衡
概述 最近.NET的世界开始闹腾了,微软官方终于加入到了对.NET跨平台的支持,并且在不久的将来,我们在VS里面写的代码可能就可以通过Mono直接在Linux和Mac上运行.那么大家(开发者和企业)为 ...
- Mysql事务探索及其在Django中的实践(二)
继上一篇<Mysql事务探索及其在Django中的实践(一)>交代完问题的背景和Mysql事务基础后,这一篇主要想介绍一下事务在Django中的使用以及实际应用给我们带来的效率提升. 首先 ...
- Mysql事务探索及其在Django中的实践(一)
前言 很早就有想开始写博客的想法,一方面是对自己近期所学知识的一些总结.沉淀,方便以后对过去的知识进行梳理.追溯,一方面也希望能通过博客来认识更多相同技术圈的朋友.所幸近期通过了博客园的申请,那么今天 ...
- netty5 HTTP协议栈浅析与实践
一.说在前面的话 前段时间,工作上需要做一个针对视频质量的统计分析系统,各端(PC端.移动端和 WEB端)将视频质量数据放在一个 HTTP 请求中上报到服务器,服务器对数据进行解析.分拣后从不同的 ...
随机推荐
- RASP Runtime Application Self-protection 运行时应用自我保护 介绍及优缺点
RASP 介绍 Runtime Application Self-protection 运行时应用自我保护 [图源:绿盟科技] 概念 Gartner (著名信息技术研究和分析厂商) 在2014年提出了 ...
- 【模型推理】量化实现分享三:详解 ACIQ 对称量化算法实现
欢迎关注我的公众号 [极智视界],回复001获取Google编程规范 O_o >_< o_O O_o ~_~ o_O 大家好,我是极智视界,本文剖析一下AC ...
- leetcode日记本
写在前面: 2019.6开始经过一年的学习,我依然没有学会算法,依然停留在最基本的阶段,面对题目依然一头雾水 但是难不是放弃的理由,根据毛主席的论持久战原理,我决定一天看一点循序渐进,相信总有一天可以 ...
- 安装rebar3
下载编译好的版本 wget https://s3.amazonaws.com/rebar3/rebar3 chmod +x /home/hylink/rebar3 (赋权) ./rebar3 loca ...
- CentOS7添加开机启动服务或脚本
方法一(rc.local) 改方式配置自动启动最为简单,只需要修改rc.local文件 由于在centos7中/etc/rc.d/rc.local的权限被降低了,所以需要赋予其可执行权 chmod + ...
- mysql 5.7.29 在centos7.6下超简单的本地yum源安装与配置
目录 生成yum源元数据 从网易镜像站下载MySQL 5.7 的 bundle包 创建文件 mysql-local.repo 执行yum install命令 生成yum源元数据 createrepo ...
- js跨域请求解决方案
什么是跨域? 跨域是指一个域下的文档或脚本试图去请求另一个域下的资源,这里跨域是广义的. 广义的跨域: 1.) 资源跳转: A链接.重定向.表单提交 2.) 资源嵌入: <link>.&l ...
- nuxt服务端渲染
<template> <div class="page"> page is search <ul> <li v-for="(it ...
- hal 编码器做用户输入时捕获初值的设置
uint16_t encoderDirection = __HAL_TIM_IS_TIM_COUNTING_DOWN(&htim3); uint16_t encoderValue = __HA ...
- 灵雀云发布云原生制品仓库Harbor企业版(Alauda Registry Service for Harbor)
灵雀云发布云原生制品仓库Harbor企业版(Alauda Registry Service for Harbor) 近日,国内领先的云原生全栈私有云提供商灵雀云宣布,推出企业版云原生制品仓库Ala ...