谷歌“史上最强GAN”,现在有了PyTorch预训练版,可直接玩耍 | 代码

耶,今天又是鬼畜的一天。

魔栗 发自 凹非寺

量子位 报道 | 公众号 QbitAI

你知道么,和BigGAN一起玩耍,会上瘾的。

比如,生成了一只狗,再生成了一只汉堡。

谷歌“史上最强GAN”,现在有了PyTorch预训练版,快来玩 | 代码

那么,狗 × 汉堡 = ?

谷歌“史上最强GAN”,现在有了PyTorch预训练版,快来玩 | 代码

一看就是亲生的。

现在,身为PyTorch用户的你,也可以拥有一只BigGAN,而且不用自己训练,便能直接玩耍。

一向以造福人类为己任的抱抱脸 (Hugging Face) 团队,用PyTorch复现了这个“史上最强”GAN。

团队开源了预训练模型,只要pip install一下,你有什么大胆想法,就可以实施了。

推特用户纷纷表示欢迎:

谷歌“史上最强GAN”,现在有了PyTorch预训练版,快来玩 | 代码

还原度极高

开源项目里有三个模型,是不同分辨率的bigGAN:

128×128,256×256,512×512。

抱抱脸团队说,模型的参数都是BigGAN的爸爸DeepMind官方训练的成果。

谷歌“史上最强GAN”,现在有了PyTorch预训练版,快来玩 | 代码

团队说,他们是用官方的原始计算图 (Computation Graph) 来复现的,与原模型的表现几乎无差:输出差异 (Output Difference) 的方差在10^-5级。

官方模型是放在TensorFlow Hub上,抱抱脸还提供了把TF模型转成PyTorch模型时,用到的脚本。

更加温柔的是,最后会显示生成效果:

谷歌“史上最强GAN”,现在有了PyTorch预训练版,快来玩 | 代码

肉眼看去,成果喜人。

食用方法

如果只是想随意玩耍的话,pip install就够了。

如果要用前面提到的转换脚本,以及ImageNet实用程序的话,就要再安装一些依赖项。记得要用full_requirements.txt来装:

1 git clone https://github.com/huggingface/pytorch-pretrained-BigGAN.git
2 cd pytorch-pretrained-BigGAN
3 pip install -r full_requirements.txt

128×128模型,有5040多万参数;256×256模型,有5590多万参数;512×512模型,有5620多万参数。三个模型,大小都在200~Mb。

安装之后,正式开始食用:

 1 import torch
 2 from pytorch_pretrained_biggan import (BigGAN, one_hot_from_names, truncated_noise_sample,
 3 save_as_images, display_in_terminal)
 4
 5 # OPTIONAL: if you want to have more information on what's happening, activate the logger as follows
 6 import logging
 7 logging.basicConfig(level=logging.INFO)
 8
 9 # Load pre-trained model tokenizer (vocabulary)
10 model = BigGAN.from_pretrained('biggan-deep-256')
11
12 # Prepare a input
13 truncation = 0.4
14 class_vector = one_hot_from_names(['soap bubble', 'coffee', 'mushroom'], batch_size=3)
15 noise_vector = truncated_noise_sample(truncation=truncation, batch_size=3)
16
17 # All in tensors
18 noise_vector = torch.from_numpy(noise_vector)
19 class_vector = torch.from_numpy(class_vector)
20
21 # If you have a GPU, put everything on cuda
22 noise_vector = noise_vector.to('cuda')
23 class_vector = class_vector.to('cuda')
24 model.to('cuda')
25
26 # Generate an image
27 with torch.no_grad():
28 output = model(noise_vector, class_vector, truncation)
29
30 # If you have a GPU put back on CPU
31 output = output.to('cpu')
32
33 # If you have a sixtel compatible terminal you can display the images in the terminal
34 # (see https://github.com/saitoha/libsixel for details)
35 display_in_terminal(output)
36
37 # Save results as png images
38 save_as_images(output)

到这里,图像就愉快地生成了。

那么,你有大胆的想法了么?

举个栗子,“红酒烩鸡”:

谷歌“史上最强GAN”,现在有了PyTorch预训练版,快来玩 | 代码

代码传送门:

https://github.com/huggingface/pytorch-pretrained-BigGAN

—  —

版权所有,未经授权不得以任何形式转载及使用,违者必究。