Showing preview only (319K chars total). Download the full file or copy to clipboard to get everything.
Repository: gnobitab/FuseDream
Branch: main
Commit: 3fe28b8db3e0
Files: 48
Total size: 303.0 KB
Directory structure:
gitextract_02nfnavz/
├── BigGAN_utils/
│ ├── BigGAN.py
│ ├── BigGANdeep.py
│ ├── LICENSE
│ ├── README.md
│ ├── TFHub/
│ │ ├── biggan_v1.py
│ │ └── converter.py
│ ├── __init__.py
│ ├── animal_hash.py
│ ├── binary_utils.py
│ ├── calculate_inception_moments.py
│ ├── datasets.py
│ ├── inception_tf13.py
│ ├── inception_utils.py
│ ├── layers.py
│ ├── logs/
│ │ ├── BigGAN_ch96_bs256x8.jsonl
│ │ ├── compare_IS.m
│ │ ├── metalog.txt
│ │ ├── process_inception_log.m
│ │ └── process_training.m
│ ├── losses.py
│ ├── make_hdf5.py
│ ├── sample.py
│ ├── scripts/
│ │ ├── launch_BigGAN_bs256x8.sh
│ │ ├── launch_BigGAN_bs512x4.sh
│ │ ├── launch_BigGAN_ch64_bs256x8.sh
│ │ ├── launch_BigGAN_deep.sh
│ │ ├── launch_SAGAN_bs128x2_ema.sh
│ │ ├── launch_SNGAN.sh
│ │ ├── launch_cifar_ema.sh
│ │ ├── sample_BigGAN_bs256x8.sh
│ │ ├── sample_cifar_ema.sh
│ │ └── utils/
│ │ ├── duplicate.sh
│ │ └── prepare_data.sh
│ ├── sync_batchnorm/
│ │ ├── __init__.py
│ │ ├── batchnorm.py
│ │ ├── batchnorm_reimpl.py
│ │ ├── comm.py
│ │ ├── replicate.py
│ │ └── unittest.py
│ ├── train.py
│ ├── train_fns.py
│ ├── utils.py
│ └── weights/
│ └── README.md
├── DiffAugment_pytorch.py
├── LICENSE
├── README.md
├── fusedream_generator.py
└── fusedream_utils.py
================================================
FILE CONTENTS
================================================
================================================
FILE: BigGAN_utils/BigGAN.py
================================================
import numpy as np
import math
import functools
import torch
import torch.nn as nn
from torch.nn import init
import torch.optim as optim
import torch.nn.functional as F
from torch.nn import Parameter as P
import layers
from sync_batchnorm import SynchronizedBatchNorm2d as SyncBatchNorm2d
# Architectures for G
# Attention is passed in in the format '32_64' to mean applying an attention
# block at both resolution 32x32 and 64x64. Just '64' will apply at 64x64.
def G_arch(ch=64, attention='64', ksize='333333', dilation='111111'):
arch = {}
arch[512] = {'in_channels' : [ch * item for item in [16, 16, 8, 8, 4, 2, 1]],
'out_channels' : [ch * item for item in [16, 8, 8, 4, 2, 1, 1]],
'upsample' : [True] * 7,
'resolution' : [8, 16, 32, 64, 128, 256, 512],
'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
for i in range(3,10)}}
arch[256] = {'in_channels' : [ch * item for item in [16, 16, 8, 8, 4, 2]],
'out_channels' : [ch * item for item in [16, 8, 8, 4, 2, 1]],
'upsample' : [True] * 6,
'resolution' : [8, 16, 32, 64, 128, 256],
'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
for i in range(3,9)}}
arch[128] = {'in_channels' : [ch * item for item in [16, 16, 8, 4, 2]],
'out_channels' : [ch * item for item in [16, 8, 4, 2, 1]],
'upsample' : [True] * 5,
'resolution' : [8, 16, 32, 64, 128],
'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
for i in range(3,8)}}
arch[64] = {'in_channels' : [ch * item for item in [16, 16, 8, 4]],
'out_channels' : [ch * item for item in [16, 8, 4, 2]],
'upsample' : [True] * 4,
'resolution' : [8, 16, 32, 64],
'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
for i in range(3,7)}}
arch[32] = {'in_channels' : [ch * item for item in [4, 4, 4]],
'out_channels' : [ch * item for item in [4, 4, 4]],
'upsample' : [True] * 3,
'resolution' : [8, 16, 32],
'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
for i in range(3,6)}}
return arch
class Generator(nn.Module):
def __init__(self, G_ch=64, dim_z=128, bottom_width=4, resolution=128,
G_kernel_size=3, G_attn='64', n_classes=1000,
num_G_SVs=1, num_G_SV_itrs=1,
G_shared=True, shared_dim=0, hier=False,
cross_replica=False, mybn=False,
G_activation=nn.ReLU(inplace=False),
G_lr=5e-5, G_B1=0.0, G_B2=0.999, adam_eps=1e-8,
BN_eps=1e-5, SN_eps=1e-12, G_mixed_precision=False, G_fp16=False,
G_init='ortho', skip_init=False, no_optim=False,
G_param='SN', norm_style='bn',
**kwargs):
super(Generator, self).__init__()
# Channel width mulitplier
self.ch = G_ch
# Dimensionality of the latent space
self.dim_z = dim_z
# The initial spatial dimensions
self.bottom_width = bottom_width
# Resolution of the output
self.resolution = resolution
# Kernel size?
self.kernel_size = G_kernel_size
# Attention?
self.attention = G_attn
# number of classes, for use in categorical conditional generation
self.n_classes = n_classes
# Use shared embeddings?
self.G_shared = G_shared
# Dimensionality of the shared embedding? Unused if not using G_shared
self.shared_dim = shared_dim if shared_dim > 0 else dim_z
# Hierarchical latent space?
self.hier = hier
# Cross replica batchnorm?
self.cross_replica = cross_replica
# Use my batchnorm?
self.mybn = mybn
# nonlinearity for residual blocks
self.activation = G_activation
# Initialization style
self.init = G_init
# Parameterization style
self.G_param = G_param
# Normalization style
self.norm_style = norm_style
# Epsilon for BatchNorm?
self.BN_eps = BN_eps
# Epsilon for Spectral Norm?
self.SN_eps = SN_eps
# fp16?
self.fp16 = G_fp16
# Architecture dict
self.arch = G_arch(self.ch, self.attention)[resolution]
# If using hierarchical latents, adjust z
if self.hier:
# Number of places z slots into
self.num_slots = len(self.arch['in_channels']) + 1
self.z_chunk_size = (self.dim_z // self.num_slots)
# Recalculate latent dimensionality for even splitting into chunks
self.dim_z = self.z_chunk_size * self.num_slots
else:
self.num_slots = 1
self.z_chunk_size = 0
# Which convs, batchnorms, and linear layers to use
if self.G_param == 'SN':
self.which_conv = functools.partial(layers.SNConv2d,
kernel_size=3, padding=1,
num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
eps=self.SN_eps)
self.which_linear = functools.partial(layers.SNLinear,
num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
eps=self.SN_eps)
else:
self.which_conv = functools.partial(nn.Conv2d, kernel_size=3, padding=1)
self.which_linear = nn.Linear
# We use a non-spectral-normed embedding here regardless;
# For some reason applying SN to G's embedding seems to randomly cripple G
self.which_embedding = nn.Embedding
bn_linear = (functools.partial(self.which_linear, bias=False) if self.G_shared
else self.which_embedding)
self.which_bn = functools.partial(layers.ccbn,
which_linear=bn_linear,
cross_replica=self.cross_replica,
mybn=self.mybn,
input_size=(self.shared_dim + self.z_chunk_size if self.G_shared
else self.n_classes),
norm_style=self.norm_style,
eps=self.BN_eps)
# Prepare model
# If not using shared embeddings, self.shared is just a passthrough
self.shared = (self.which_embedding(n_classes, self.shared_dim) if G_shared
else layers.identity())
# First linear layer
self.linear = self.which_linear(self.dim_z // self.num_slots,
self.arch['in_channels'][0] * (self.bottom_width **2))
# self.blocks is a doubly-nested list of modules, the outer loop intended
# to be over blocks at a given resolution (resblocks and/or self-attention)
# while the inner loop is over a given block
self.blocks = []
for index in range(len(self.arch['out_channels'])):
self.blocks += [[layers.GBlock(in_channels=self.arch['in_channels'][index],
out_channels=self.arch['out_channels'][index],
which_conv=self.which_conv,
which_bn=self.which_bn,
activation=self.activation,
upsample=(functools.partial(F.interpolate, scale_factor=2)
if self.arch['upsample'][index] else None))]]
# If attention on this block, attach it to the end
if self.arch['attention'][self.arch['resolution'][index]]:
print(self.arch['resolution'], self.arch['attention'])
print('Adding attention layer in G at resolution %d' % self.arch['resolution'][index])
self.blocks[-1] += [layers.Attention(self.arch['out_channels'][index], self.which_conv)]
# Turn self.blocks into a ModuleList so that it's all properly registered.
self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])
# output layer: batchnorm-relu-conv.
# Consider using a non-spectral conv here
self.output_layer = nn.Sequential(layers.bn(self.arch['out_channels'][-1],
cross_replica=self.cross_replica,
mybn=self.mybn),
self.activation,
self.which_conv(self.arch['out_channels'][-1], 3))
# Initialize weights. Optionally skip init for testing.
if not skip_init:
self.init_weights()
# Set up optimizer
# If this is an EMA copy, no need for an optim, so just return now
if no_optim:
return
self.lr, self.B1, self.B2, self.adam_eps = G_lr, G_B1, G_B2, adam_eps
if G_mixed_precision:
print('Using fp16 adam in G...')
import utils
self.optim = utils.Adam16(params=self.parameters(), lr=self.lr,
betas=(self.B1, self.B2), weight_decay=0,
eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(), lr=self.lr,
betas=(self.B1, self.B2), weight_decay=0,
eps=self.adam_eps)
# LR scheduling, left here for forward compatibility
# self.lr_sched = {'itr' : 0}# if self.progressive else {}
# self.j = 0
# Initialize
def init_weights(self):
self.param_count = 0
for module in self.modules():
if (isinstance(module, nn.Conv2d)
or isinstance(module, nn.Linear)
or isinstance(module, nn.Embedding)):
if self.init == 'ortho':
init.orthogonal_(module.weight)
elif self.init == 'N02':
init.normal_(module.weight, 0, 0.02)
elif self.init in ['glorot', 'xavier']:
init.xavier_uniform_(module.weight)
else:
print('Init style not recognized...')
self.param_count += sum([p.data.nelement() for p in module.parameters()])
print('Param count for G''s initialized parameters: %d' % self.param_count)
# Note on this forward function: we pass in a y vector which has
# already been passed through G.shared to enable easy class-wise
# interpolation later. If we passed in the one-hot and then ran it through
# G.shared in this forward function, it would be harder to handle.
def forward(self, z, y, w_y=None):
if w_y is not None:
s_y = torch.softmax(w_y, dim=1)
cur_y = s_y * y
y = cur_y.sum(dim=1, keepdim=False)
# If hierarchical, concatenate zs and ys
if self.hier:
zs = torch.split(z, self.z_chunk_size, 1)
z = zs[0]
ys = [torch.cat([y, item], 1) for item in zs[1:]]
else:
ys = [y] * len(self.blocks)
# First linear layer
h = self.linear(z)
# Reshape
h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width)
# Loop over blocks
for index, blocklist in enumerate(self.blocks):
# Second inner loop in case block has multiple layers
for block in blocklist:
h = block(h, ys[index])
# Apply batchnorm-relu-conv-tanh at output
return torch.tanh(self.output_layer(h))
# Note on this forward function: we pass in a y vector which has
# already been passed through G.shared to enable easy class-wise
# interpolation later. If we passed in the one-hot and then ran it through
# G.shared in this forward function, it would be harder to handle.
def forward_org(self, z, y):
# If hierarchical, concatenate zs and ys
if self.hier:
zs = torch.split(z, self.z_chunk_size, 1)
z = zs[0]
ys = [torch.cat([y, item], 1) for item in zs[1:]]
else:
ys = [y] * len(self.blocks)
# First linear layer
h = self.linear(z)
# Reshape
h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width)
# Loop over blocks
for index, blocklist in enumerate(self.blocks):
# Second inner loop in case block has multiple layers
for block in blocklist:
h = block(h, ys[index])
# Apply batchnorm-relu-conv-tanh at output
return torch.tanh(self.output_layer(h))
# Discriminator architecture, same paradigm as G's above
def D_arch(ch=64, attention='64',ksize='333333', dilation='111111'):
arch = {}
arch[256] = {'in_channels' : [3] + [ch*item for item in [1, 2, 4, 8, 8, 16]],
'out_channels' : [item * ch for item in [1, 2, 4, 8, 8, 16, 16]],
'downsample' : [True] * 6 + [False],
'resolution' : [128, 64, 32, 16, 8, 4, 4 ],
'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
for i in range(2,8)}}
arch[128] = {'in_channels' : [3] + [ch*item for item in [1, 2, 4, 8, 16]],
'out_channels' : [item * ch for item in [1, 2, 4, 8, 16, 16]],
'downsample' : [True] * 5 + [False],
'resolution' : [64, 32, 16, 8, 4, 4],
'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
for i in range(2,8)}}
arch[64] = {'in_channels' : [3] + [ch*item for item in [1, 2, 4, 8]],
'out_channels' : [item * ch for item in [1, 2, 4, 8, 16]],
'downsample' : [True] * 4 + [False],
'resolution' : [32, 16, 8, 4, 4],
'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
for i in range(2,7)}}
arch[32] = {'in_channels' : [3] + [item * ch for item in [4, 4, 4]],
'out_channels' : [item * ch for item in [4, 4, 4, 4]],
'downsample' : [True, True, False, False],
'resolution' : [16, 16, 16, 16],
'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
for i in range(2,6)}}
return arch
class Discriminator(nn.Module):
def __init__(self, D_ch=64, D_wide=True, resolution=128,
D_kernel_size=3, D_attn='64', n_classes=1000,
num_D_SVs=1, num_D_SV_itrs=1, D_activation=nn.ReLU(inplace=False),
D_lr=2e-4, D_B1=0.0, D_B2=0.999, adam_eps=1e-8,
SN_eps=1e-12, output_dim=1, D_mixed_precision=False, D_fp16=False,
D_init='ortho', skip_init=False, D_param='SN', **kwargs):
super(Discriminator, self).__init__()
# Width multiplier
self.ch = D_ch
# Use Wide D as in BigGAN and SA-GAN or skinny D as in SN-GAN?
self.D_wide = D_wide
# Resolution
self.resolution = resolution
# Kernel size
self.kernel_size = D_kernel_size
# Attention?
self.attention = D_attn
# Number of classes
self.n_classes = n_classes
# Activation
self.activation = D_activation
# Initialization style
self.init = D_init
# Parameterization style
self.D_param = D_param
# Epsilon for Spectral Norm?
self.SN_eps = SN_eps
# Fp16?
self.fp16 = D_fp16
# Architecture
self.arch = D_arch(self.ch, self.attention)[resolution]
# Which convs, batchnorms, and linear layers to use
# No option to turn off SN in D right now
if self.D_param == 'SN':
self.which_conv = functools.partial(layers.SNConv2d,
kernel_size=3, padding=1,
num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
self.which_linear = functools.partial(layers.SNLinear,
num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
self.which_embedding = functools.partial(layers.SNEmbedding,
num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
# Prepare model
# self.blocks is a doubly-nested list of modules, the outer loop intended
# to be over blocks at a given resolution (resblocks and/or self-attention)
self.blocks = []
for index in range(len(self.arch['out_channels'])):
self.blocks += [[layers.DBlock(in_channels=self.arch['in_channels'][index],
out_channels=self.arch['out_channels'][index],
which_conv=self.which_conv,
wide=self.D_wide,
activation=self.activation,
preactivation=(index > 0),
downsample=(nn.AvgPool2d(2) if self.arch['downsample'][index] else None))]]
# If attention on this block, attach it to the end
if self.arch['attention'][self.arch['resolution'][index]]:
print('Adding attention layer in D at resolution %d' % self.arch['resolution'][index])
self.blocks[-1] += [layers.Attention(self.arch['out_channels'][index],
self.which_conv)]
# Turn self.blocks into a ModuleList so that it's all properly registered.
self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])
# Linear output layer. The output dimension is typically 1, but may be
# larger if we're e.g. turning this into a VAE with an inference output
self.linear = self.which_linear(self.arch['out_channels'][-1], output_dim)
# Embedding for projection discrimination
self.embed = self.which_embedding(self.n_classes, self.arch['out_channels'][-1])
# Initialize weights
if not skip_init:
self.init_weights()
# Set up optimizer
self.lr, self.B1, self.B2, self.adam_eps = D_lr, D_B1, D_B2, adam_eps
if D_mixed_precision:
print('Using fp16 adam in D...')
import utils
self.optim = utils.Adam16(params=self.parameters(), lr=self.lr,
betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(), lr=self.lr,
betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)
# LR scheduling, left here for forward compatibility
# self.lr_sched = {'itr' : 0}# if self.progressive else {}
# self.j = 0
# Initialize
def init_weights(self):
self.param_count = 0
for module in self.modules():
if (isinstance(module, nn.Conv2d)
or isinstance(module, nn.Linear)
or isinstance(module, nn.Embedding)):
if self.init == 'ortho':
init.orthogonal_(module.weight)
elif self.init == 'N02':
init.normal_(module.weight, 0, 0.02)
elif self.init in ['glorot', 'xavier']:
init.xavier_uniform_(module.weight)
else:
print('Init style not recognized...')
self.param_count += sum([p.data.nelement() for p in module.parameters()])
print('Param count for D''s initialized parameters: %d' % self.param_count)
def forward(self, x, y=None):
# Stick x into h for cleaner for loops without flow control
h = x
# Loop over blocks
for index, blocklist in enumerate(self.blocks):
for block in blocklist:
h = block(h)
# Apply global sum pooling as in SN-GAN
h = torch.sum(self.activation(h), [2, 3])
# Get initial class-unconditional output
out = self.linear(h)
# Get projection of final featureset onto class vectors and add to evidence
out = out + torch.sum(self.embed(y) * h, 1, keepdim=True)
return out
# Parallelized G_D to minimize cross-gpu communication
# Without this, Generator outputs would get all-gathered and then rebroadcast.
class G_D(nn.Module):
def __init__(self, G, D):
super(G_D, self).__init__()
self.G = G
self.D = D
def forward(self, z, gy, x=None, dy=None, train_G=False, return_G_z=False,
split_D=False):
# If training G, enable grad tape
with torch.set_grad_enabled(train_G):
# Get Generator output given noise
G_z = self.G(z, self.G.shared(gy))
# Cast as necessary
if self.G.fp16 and not self.D.fp16:
G_z = G_z.float()
if self.D.fp16 and not self.G.fp16:
G_z = G_z.half()
# Split_D means to run D once with real data and once with fake,
# rather than concatenating along the batch dimension.
if split_D:
D_fake = self.D(G_z, gy)
if x is not None:
D_real = self.D(x, dy)
return D_fake, D_real
else:
if return_G_z:
return D_fake, G_z
else:
return D_fake
# If real data is provided, concatenate it with the Generator's output
# along the batch dimension for improved efficiency.
else:
D_input = torch.cat([G_z, x], 0) if x is not None else G_z
D_class = torch.cat([gy, dy], 0) if dy is not None else gy
# Get Discriminator output
D_out = self.D(D_input, D_class)
if x is not None:
return torch.split(D_out, [G_z.shape[0], x.shape[0]]) # D_fake, D_real
else:
if return_G_z:
return D_out, G_z
else:
return D_out
================================================
FILE: BigGAN_utils/BigGANdeep.py
================================================
import numpy as np
import math
import functools
import torch
import torch.nn as nn
from torch.nn import init
import torch.optim as optim
import torch.nn.functional as F
from torch.nn import Parameter as P
import layers
from sync_batchnorm import SynchronizedBatchNorm2d as SyncBatchNorm2d
# BigGAN-deep: uses a different resblock and pattern
# Architectures for G
# Attention is passed in in the format '32_64' to mean applying an attention
# block at both resolution 32x32 and 64x64. Just '64' will apply at 64x64.
# Channel ratio is the ratio of
class GBlock(nn.Module):
def __init__(self, in_channels, out_channels,
which_conv=nn.Conv2d, which_bn=layers.bn, activation=None,
upsample=None, channel_ratio=4):
super(GBlock, self).__init__()
self.in_channels, self.out_channels = in_channels, out_channels
self.hidden_channels = self.in_channels // channel_ratio
self.which_conv, self.which_bn = which_conv, which_bn
self.activation = activation
# Conv layers
self.conv1 = self.which_conv(self.in_channels, self.hidden_channels,
kernel_size=1, padding=0)
self.conv2 = self.which_conv(self.hidden_channels, self.hidden_channels)
self.conv3 = self.which_conv(self.hidden_channels, self.hidden_channels)
self.conv4 = self.which_conv(self.hidden_channels, self.out_channels,
kernel_size=1, padding=0)
# Batchnorm layers
self.bn1 = self.which_bn(self.in_channels)
self.bn2 = self.which_bn(self.hidden_channels)
self.bn3 = self.which_bn(self.hidden_channels)
self.bn4 = self.which_bn(self.hidden_channels)
# upsample layers
self.upsample = upsample
def forward(self, x, y):
# Project down to channel ratio
h = self.conv1(self.activation(self.bn1(x, y)))
# Apply next BN-ReLU
h = self.activation(self.bn2(h, y))
# Drop channels in x if necessary
if self.in_channels != self.out_channels:
x = x[:, :self.out_channels]
# Upsample both h and x at this point
if self.upsample:
h = self.upsample(h)
x = self.upsample(x)
# 3x3 convs
h = self.conv2(h)
h = self.conv3(self.activation(self.bn3(h, y)))
# Final 1x1 conv
h = self.conv4(self.activation(self.bn4(h, y)))
return h + x
def G_arch(ch=64, attention='64', ksize='333333', dilation='111111'):
arch = {}
arch[256] = {'in_channels' : [ch * item for item in [16, 16, 8, 8, 4, 2]],
'out_channels' : [ch * item for item in [16, 8, 8, 4, 2, 1]],
'upsample' : [True] * 6,
'resolution' : [8, 16, 32, 64, 128, 256],
'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
for i in range(3,9)}}
arch[128] = {'in_channels' : [ch * item for item in [16, 16, 8, 4, 2]],
'out_channels' : [ch * item for item in [16, 8, 4, 2, 1]],
'upsample' : [True] * 5,
'resolution' : [8, 16, 32, 64, 128],
'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
for i in range(3,8)}}
arch[64] = {'in_channels' : [ch * item for item in [16, 16, 8, 4]],
'out_channels' : [ch * item for item in [16, 8, 4, 2]],
'upsample' : [True] * 4,
'resolution' : [8, 16, 32, 64],
'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
for i in range(3,7)}}
arch[32] = {'in_channels' : [ch * item for item in [4, 4, 4]],
'out_channels' : [ch * item for item in [4, 4, 4]],
'upsample' : [True] * 3,
'resolution' : [8, 16, 32],
'attention' : {2**i: (2**i in [int(item) for item in attention.split('_')])
for i in range(3,6)}}
return arch
class Generator(nn.Module):
def __init__(self, G_ch=64, G_depth=2, dim_z=128, bottom_width=4, resolution=128,
G_kernel_size=3, G_attn='64', n_classes=1000,
num_G_SVs=1, num_G_SV_itrs=1,
G_shared=True, shared_dim=0, hier=False,
cross_replica=False, mybn=False,
G_activation=nn.ReLU(inplace=False),
G_lr=5e-5, G_B1=0.0, G_B2=0.999, adam_eps=1e-8,
BN_eps=1e-5, SN_eps=1e-12, G_mixed_precision=False, G_fp16=False,
G_init='ortho', skip_init=False, no_optim=False,
G_param='SN', norm_style='bn',
**kwargs):
super(Generator, self).__init__()
# Channel width mulitplier
self.ch = G_ch
# Number of resblocks per stage
self.G_depth = G_depth
# Dimensionality of the latent space
self.dim_z = dim_z
# The initial spatial dimensions
self.bottom_width = bottom_width
# Resolution of the output
self.resolution = resolution
# Kernel size?
self.kernel_size = G_kernel_size
# Attention?
self.attention = G_attn
# number of classes, for use in categorical conditional generation
self.n_classes = n_classes
# Use shared embeddings?
self.G_shared = G_shared
# Dimensionality of the shared embedding? Unused if not using G_shared
self.shared_dim = shared_dim if shared_dim > 0 else dim_z
# Hierarchical latent space?
self.hier = hier
# Cross replica batchnorm?
self.cross_replica = cross_replica
# Use my batchnorm?
self.mybn = mybn
# nonlinearity for residual blocks
self.activation = G_activation
# Initialization style
self.init = G_init
# Parameterization style
self.G_param = G_param
# Normalization style
self.norm_style = norm_style
# Epsilon for BatchNorm?
self.BN_eps = BN_eps
# Epsilon for Spectral Norm?
self.SN_eps = SN_eps
# fp16?
self.fp16 = G_fp16
# Architecture dict
self.arch = G_arch(self.ch, self.attention)[resolution]
# Which convs, batchnorms, and linear layers to use
if self.G_param == 'SN':
self.which_conv = functools.partial(layers.SNConv2d,
kernel_size=3, padding=1,
num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
eps=self.SN_eps)
self.which_linear = functools.partial(layers.SNLinear,
num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
eps=self.SN_eps)
else:
self.which_conv = functools.partial(nn.Conv2d, kernel_size=3, padding=1)
self.which_linear = nn.Linear
# We use a non-spectral-normed embedding here regardless;
# For some reason applying SN to G's embedding seems to randomly cripple G
self.which_embedding = nn.Embedding
bn_linear = (functools.partial(self.which_linear, bias=False) if self.G_shared
else self.which_embedding)
self.which_bn = functools.partial(layers.ccbn,
which_linear=bn_linear,
cross_replica=self.cross_replica,
mybn=self.mybn,
input_size=(self.shared_dim + self.dim_z if self.G_shared
else self.n_classes),
norm_style=self.norm_style,
eps=self.BN_eps)
# Prepare model
# If not using shared embeddings, self.shared is just a passthrough
self.shared = (self.which_embedding(n_classes, self.shared_dim) if G_shared
else layers.identity())
# First linear layer
self.linear = self.which_linear(self.dim_z + self.shared_dim, self.arch['in_channels'][0] * (self.bottom_width **2))
# self.blocks is a doubly-nested list of modules, the outer loop intended
# to be over blocks at a given resolution (resblocks and/or self-attention)
# while the inner loop is over a given block
self.blocks = []
for index in range(len(self.arch['out_channels'])):
self.blocks += [[GBlock(in_channels=self.arch['in_channels'][index],
out_channels=self.arch['in_channels'][index] if g_index==0 else self.arch['out_channels'][index],
which_conv=self.which_conv,
which_bn=self.which_bn,
activation=self.activation,
upsample=(functools.partial(F.interpolate, scale_factor=2)
if self.arch['upsample'][index] and g_index == (self.G_depth-1) else None))]
for g_index in range(self.G_depth)]
# If attention on this block, attach it to the end
if self.arch['attention'][self.arch['resolution'][index]]:
print('Adding attention layer in G at resolution %d' % self.arch['resolution'][index])
self.blocks[-1] += [layers.Attention(self.arch['out_channels'][index], self.which_conv)]
# Turn self.blocks into a ModuleList so that it's all properly registered.
self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])
# output layer: batchnorm-relu-conv.
# Consider using a non-spectral conv here
self.output_layer = nn.Sequential(layers.bn(self.arch['out_channels'][-1],
cross_replica=self.cross_replica,
mybn=self.mybn),
self.activation,
self.which_conv(self.arch['out_channels'][-1], 3))
# Initialize weights. Optionally skip init for testing.
if not skip_init:
self.init_weights()
# Set up optimizer
# If this is an EMA copy, no need for an optim, so just return now
if no_optim:
return
self.lr, self.B1, self.B2, self.adam_eps = G_lr, G_B1, G_B2, adam_eps
if G_mixed_precision:
print('Using fp16 adam in G...')
import utils
self.optim = utils.Adam16(params=self.parameters(), lr=self.lr,
betas=(self.B1, self.B2), weight_decay=0,
eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(), lr=self.lr,
betas=(self.B1, self.B2), weight_decay=0,
eps=self.adam_eps)
# LR scheduling, left here for forward compatibility
# self.lr_sched = {'itr' : 0}# if self.progressive else {}
# self.j = 0
# Initialize
def init_weights(self):
self.param_count = 0
for module in self.modules():
if (isinstance(module, nn.Conv2d)
or isinstance(module, nn.Linear)
or isinstance(module, nn.Embedding)):
if self.init == 'ortho':
init.orthogonal_(module.weight)
elif self.init == 'N02':
init.normal_(module.weight, 0, 0.02)
elif self.init in ['glorot', 'xavier']:
init.xavier_uniform_(module.weight)
else:
print('Init style not recognized...')
self.param_count += sum([p.data.nelement() for p in module.parameters()])
print('Param count for G''s initialized parameters: %d' % self.param_count)
# Note on this forward function: we pass in a y vector which has
# already been passed through G.shared to enable easy class-wise
# interpolation later. If we passed in the one-hot and then ran it through
# G.shared in this forward function, it would be harder to handle.
# NOTE: The z vs y dichotomy here is for compatibility with not-y
def forward(self, z, y):
# If hierarchical, concatenate zs and ys
if self.hier:
z = torch.cat([y, z], 1)
y = z
# First linear layer
h = self.linear(z)
# Reshape
h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width)
# Loop over blocks
for index, blocklist in enumerate(self.blocks):
# Second inner loop in case block has multiple layers
for block in blocklist:
h = block(h, y)
# Apply batchnorm-relu-conv-tanh at output
return torch.tanh(self.output_layer(h))
class DBlock(nn.Module):
def __init__(self, in_channels, out_channels, which_conv=layers.SNConv2d, wide=True,
preactivation=True, activation=None, downsample=None,
channel_ratio=4):
super(DBlock, self).__init__()
self.in_channels, self.out_channels = in_channels, out_channels
# If using wide D (as in SA-GAN and BigGAN), change the channel pattern
self.hidden_channels = self.out_channels // channel_ratio
self.which_conv = which_conv
self.preactivation = preactivation
self.activation = activation
self.downsample = downsample
# Conv layers
self.conv1 = self.which_conv(self.in_channels, self.hidden_channels,
kernel_size=1, padding=0)
self.conv2 = self.which_conv(self.hidden_channels, self.hidden_channels)
self.conv3 = self.which_conv(self.hidden_channels, self.hidden_channels)
self.conv4 = self.which_conv(self.hidden_channels, self.out_channels,
kernel_size=1, padding=0)
self.learnable_sc = True if (in_channels != out_channels) else False
if self.learnable_sc:
self.conv_sc = self.which_conv(in_channels, out_channels - in_channels,
kernel_size=1, padding=0)
def shortcut(self, x):
if self.downsample:
x = self.downsample(x)
if self.learnable_sc:
x = torch.cat([x, self.conv_sc(x)], 1)
return x
def forward(self, x):
# 1x1 bottleneck conv
h = self.conv1(F.relu(x))
# 3x3 convs
h = self.conv2(self.activation(h))
h = self.conv3(self.activation(h))
# relu before downsample
h = self.activation(h)
# downsample
if self.downsample:
h = self.downsample(h)
# final 1x1 conv
h = self.conv4(h)
return h + self.shortcut(x)
# Discriminator architecture, same paradigm as G's above
def D_arch(ch=64, attention='64',ksize='333333', dilation='111111'):
arch = {}
arch[256] = {'in_channels' : [item * ch for item in [1, 2, 4, 8, 8, 16]],
'out_channels' : [item * ch for item in [2, 4, 8, 8, 16, 16]],
'downsample' : [True] * 6 + [False],
'resolution' : [128, 64, 32, 16, 8, 4, 4 ],
'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
for i in range(2,8)}}
arch[128] = {'in_channels' : [item * ch for item in [1, 2, 4, 8, 16]],
'out_channels' : [item * ch for item in [2, 4, 8, 16, 16]],
'downsample' : [True] * 5 + [False],
'resolution' : [64, 32, 16, 8, 4, 4],
'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
for i in range(2,8)}}
arch[64] = {'in_channels' : [item * ch for item in [1, 2, 4, 8]],
'out_channels' : [item * ch for item in [2, 4, 8, 16]],
'downsample' : [True] * 4 + [False],
'resolution' : [32, 16, 8, 4, 4],
'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
for i in range(2,7)}}
arch[32] = {'in_channels' : [item * ch for item in [4, 4, 4]],
'out_channels' : [item * ch for item in [4, 4, 4]],
'downsample' : [True, True, False, False],
'resolution' : [16, 16, 16, 16],
'attention' : {2**i: 2**i in [int(item) for item in attention.split('_')]
for i in range(2,6)}}
return arch
class Discriminator(nn.Module):
def __init__(self, D_ch=64, D_wide=True, D_depth=2, resolution=128,
D_kernel_size=3, D_attn='64', n_classes=1000,
num_D_SVs=1, num_D_SV_itrs=1, D_activation=nn.ReLU(inplace=False),
D_lr=2e-4, D_B1=0.0, D_B2=0.999, adam_eps=1e-8,
SN_eps=1e-12, output_dim=1, D_mixed_precision=False, D_fp16=False,
D_init='ortho', skip_init=False, D_param='SN', **kwargs):
super(Discriminator, self).__init__()
# Width multiplier
self.ch = D_ch
# Use Wide D as in BigGAN and SA-GAN or skinny D as in SN-GAN?
self.D_wide = D_wide
# How many resblocks per stage?
self.D_depth = D_depth
# Resolution
self.resolution = resolution
# Kernel size
self.kernel_size = D_kernel_size
# Attention?
self.attention = D_attn
# Number of classes
self.n_classes = n_classes
# Activation
self.activation = D_activation
# Initialization style
self.init = D_init
# Parameterization style
self.D_param = D_param
# Epsilon for Spectral Norm?
self.SN_eps = SN_eps
# Fp16?
self.fp16 = D_fp16
# Architecture
self.arch = D_arch(self.ch, self.attention)[resolution]
# Which convs, batchnorms, and linear layers to use
# No option to turn off SN in D right now
if self.D_param == 'SN':
self.which_conv = functools.partial(layers.SNConv2d,
kernel_size=3, padding=1,
num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
self.which_linear = functools.partial(layers.SNLinear,
num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
self.which_embedding = functools.partial(layers.SNEmbedding,
num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
# Prepare model
# Stem convolution
self.input_conv = self.which_conv(3, self.arch['in_channels'][0])
# self.blocks is a doubly-nested list of modules, the outer loop intended
# to be over blocks at a given resolution (resblocks and/or self-attention)
self.blocks = []
for index in range(len(self.arch['out_channels'])):
self.blocks += [[DBlock(in_channels=self.arch['in_channels'][index] if d_index==0 else self.arch['out_channels'][index],
out_channels=self.arch['out_channels'][index],
which_conv=self.which_conv,
wide=self.D_wide,
activation=self.activation,
preactivation=True,
downsample=(nn.AvgPool2d(2) if self.arch['downsample'][index] and d_index==0 else None))
for d_index in range(self.D_depth)]]
# If attention on this block, attach it to the end
if self.arch['attention'][self.arch['resolution'][index]]:
print('Adding attention layer in D at resolution %d' % self.arch['resolution'][index])
self.blocks[-1] += [layers.Attention(self.arch['out_channels'][index],
self.which_conv)]
# Turn self.blocks into a ModuleList so that it's all properly registered.
self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])
# Linear output layer. The output dimension is typically 1, but may be
# larger if we're e.g. turning this into a VAE with an inference output
self.linear = self.which_linear(self.arch['out_channels'][-1], output_dim)
# Embedding for projection discrimination
self.embed = self.which_embedding(self.n_classes, self.arch['out_channels'][-1])
# Initialize weights
if not skip_init:
self.init_weights()
# Set up optimizer
self.lr, self.B1, self.B2, self.adam_eps = D_lr, D_B1, D_B2, adam_eps
if D_mixed_precision:
print('Using fp16 adam in D...')
import utils
self.optim = utils.Adam16(params=self.parameters(), lr=self.lr,
betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)
else:
self.optim = optim.Adam(params=self.parameters(), lr=self.lr,
betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)
# LR scheduling, left here for forward compatibility
# self.lr_sched = {'itr' : 0}# if self.progressive else {}
# self.j = 0
# Initialize
def init_weights(self):
self.param_count = 0
for module in self.modules():
if (isinstance(module, nn.Conv2d)
or isinstance(module, nn.Linear)
or isinstance(module, nn.Embedding)):
if self.init == 'ortho':
init.orthogonal_(module.weight)
elif self.init == 'N02':
init.normal_(module.weight, 0, 0.02)
elif self.init in ['glorot', 'xavier']:
init.xavier_uniform_(module.weight)
else:
print('Init style not recognized...')
self.param_count += sum([p.data.nelement() for p in module.parameters()])
print('Param count for D''s initialized parameters: %d' % self.param_count)
def forward(self, x, y=None):
# Run input conv
h = self.input_conv(x)
# Loop over blocks
for index, blocklist in enumerate(self.blocks):
for block in blocklist:
h = block(h)
# Apply global sum pooling as in SN-GAN
h = torch.sum(self.activation(h), [2, 3])
# Get initial class-unconditional output
out = self.linear(h)
# Get projection of final featureset onto class vectors and add to evidence
out = out + torch.sum(self.embed(y) * h, 1, keepdim=True)
return out
# Parallelized G_D to minimize cross-gpu communication
# Without this, Generator outputs would get all-gathered and then rebroadcast.
class G_D(nn.Module):
def __init__(self, G, D):
super(G_D, self).__init__()
self.G = G
self.D = D
def forward(self, z, gy, x=None, dy=None, train_G=False, return_G_z=False,
split_D=False):
# If training G, enable grad tape
with torch.set_grad_enabled(train_G):
# Get Generator output given noise
G_z = self.G(z, self.G.shared(gy))
# Cast as necessary
if self.G.fp16 and not self.D.fp16:
G_z = G_z.float()
if self.D.fp16 and not self.G.fp16:
G_z = G_z.half()
# Split_D means to run D once with real data and once with fake,
# rather than concatenating along the batch dimension.
if split_D:
D_fake = self.D(G_z, gy)
if x is not None:
D_real = self.D(x, dy)
return D_fake, D_real
else:
if return_G_z:
return D_fake, G_z
else:
return D_fake
# If real data is provided, concatenate it with the Generator's output
# along the batch dimension for improved efficiency.
else:
D_input = torch.cat([G_z, x], 0) if x is not None else G_z
D_class = torch.cat([gy, dy], 0) if dy is not None else gy
# Get Discriminator output
D_out = self.D(D_input, D_class)
if x is not None:
return torch.split(D_out, [G_z.shape[0], x.shape[0]]) # D_fake, D_real
else:
if return_G_z:
return D_out, G_z
else:
return D_out
================================================
FILE: BigGAN_utils/LICENSE
================================================
MIT License
Copyright (c) 2019 Andy Brock
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
================================================
FILE: BigGAN_utils/README.md
================================================
Download pre-trained weights from (https://drive.google.com/drive/folders/1nJ3HmgYgeA9NZr-oU-enqbYeO7zBaANs?usp=sharing) and put them in `./weights/`
================================================
FILE: BigGAN_utils/TFHub/biggan_v1.py
================================================
# BigGAN V1:
# This is now deprecated code used for porting the TFHub modules to pytorch,
# included here for reference only.
import numpy as np
import torch
from scipy.stats import truncnorm
from torch import nn
from torch.nn import Parameter
from torch.nn import functional as F
def l2normalize(v, eps=1e-4):
return v / (v.norm() + eps)
def truncated_z_sample(batch_size, z_dim, truncation=0.5, seed=None):
state = None if seed is None else np.random.RandomState(seed)
values = truncnorm.rvs(-2, 2, size=(batch_size, z_dim), random_state=state)
return truncation * values
def denorm(x):
out = (x + 1) / 2
return out.clamp_(0, 1)
class SpectralNorm(nn.Module):
def __init__(self, module, name='weight', power_iterations=1):
super(SpectralNorm, self).__init__()
self.module = module
self.name = name
self.power_iterations = power_iterations
if not self._made_params():
self._make_params()
def _update_u_v(self):
u = getattr(self.module, self.name + "_u")
v = getattr(self.module, self.name + "_v")
w = getattr(self.module, self.name + "_bar")
height = w.data.shape[0]
_w = w.view(height, -1)
for _ in range(self.power_iterations):
v = l2normalize(torch.matmul(_w.t(), u))
u = l2normalize(torch.matmul(_w, v))
sigma = u.dot((_w).mv(v))
setattr(self.module, self.name, w / sigma.expand_as(w))
def _made_params(self):
try:
getattr(self.module, self.name + "_u")
getattr(self.module, self.name + "_v")
getattr(self.module, self.name + "_bar")
return True
except AttributeError:
return False
def _make_params(self):
w = getattr(self.module, self.name)
height = w.data.shape[0]
width = w.view(height, -1).data.shape[1]
u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
v = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)
u.data = l2normalize(u.data)
v.data = l2normalize(v.data)
w_bar = Parameter(w.data)
del self.module._parameters[self.name]
self.module.register_parameter(self.name + "_u", u)
self.module.register_parameter(self.name + "_v", v)
self.module.register_parameter(self.name + "_bar", w_bar)
def forward(self, *args):
self._update_u_v()
return self.module.forward(*args)
class SelfAttention(nn.Module):
""" Self Attention Layer"""
def __init__(self, in_dim, activation=F.relu):
super().__init__()
self.chanel_in = in_dim
self.activation = activation
self.theta = SpectralNorm(nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1, bias=False))
self.phi = SpectralNorm(nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 8, kernel_size=1, bias=False))
self.pool = nn.MaxPool2d(2, 2)
self.g = SpectralNorm(nn.Conv2d(in_channels=in_dim, out_channels=in_dim // 2, kernel_size=1, bias=False))
self.o_conv = SpectralNorm(nn.Conv2d(in_channels=in_dim // 2, out_channels=in_dim, kernel_size=1, bias=False))
self.gamma = nn.Parameter(torch.zeros(1))
self.softmax = nn.Softmax(dim=-1)
def forward(self, x):
m_batchsize, C, width, height = x.size()
N = height * width
theta = self.theta(x)
phi = self.phi(x)
phi = self.pool(phi)
phi = phi.view(m_batchsize, -1, N // 4)
theta = theta.view(m_batchsize, -1, N)
theta = theta.permute(0, 2, 1)
attention = self.softmax(torch.bmm(theta, phi))
g = self.pool(self.g(x)).view(m_batchsize, -1, N // 4)
attn_g = torch.bmm(g, attention.permute(0, 2, 1)).view(m_batchsize, -1, width, height)
out = self.o_conv(attn_g)
return self.gamma * out + x
class ConditionalBatchNorm2d(nn.Module):
def __init__(self, num_features, num_classes, eps=1e-4, momentum=0.1):
super().__init__()
self.num_features = num_features
self.bn = nn.BatchNorm2d(num_features, affine=False, eps=eps, momentum=momentum)
self.gamma_embed = SpectralNorm(nn.Linear(num_classes, num_features, bias=False))
self.beta_embed = SpectralNorm(nn.Linear(num_classes, num_features, bias=False))
def forward(self, x, y):
out = self.bn(x)
gamma = self.gamma_embed(y) + 1
beta = self.beta_embed(y)
out = gamma.view(-1, self.num_features, 1, 1) * out + beta.view(-1, self.num_features, 1, 1)
return out
class GBlock(nn.Module):
def __init__(
self,
in_channel,
out_channel,
kernel_size=[3, 3],
padding=1,
stride=1,
n_class=None,
bn=True,
activation=F.relu,
upsample=True,
downsample=False,
z_dim=148,
):
super().__init__()
self.conv0 = SpectralNorm(
nn.Conv2d(in_channel, out_channel, kernel_size, stride, padding, bias=True if bn else True)
)
self.conv1 = SpectralNorm(
nn.Conv2d(out_channel, out_channel, kernel_size, stride, padding, bias=True if bn else True)
)
self.skip_proj = False
if in_channel != out_channel or upsample or downsample:
self.conv_sc = SpectralNorm(nn.Conv2d(in_channel, out_channel, 1, 1, 0))
self.skip_proj = True
self.upsample = upsample
self.downsample = downsample
self.activation = activation
self.bn = bn
if bn:
self.HyperBN = ConditionalBatchNorm2d(in_channel, z_dim)
self.HyperBN_1 = ConditionalBatchNorm2d(out_channel, z_dim)
def forward(self, input, condition=None):
out = input
if self.bn:
out = self.HyperBN(out, condition)
out = self.activation(out)
if self.upsample:
out = F.interpolate(out, scale_factor=2)
out = self.conv0(out)
if self.bn:
out = self.HyperBN_1(out, condition)
out = self.activation(out)
out = self.conv1(out)
if self.downsample:
out = F.avg_pool2d(out, 2)
if self.skip_proj:
skip = input
if self.upsample:
skip = F.interpolate(skip, scale_factor=2)
skip = self.conv_sc(skip)
if self.downsample:
skip = F.avg_pool2d(skip, 2)
else:
skip = input
return out + skip
class Generator128(nn.Module):
def __init__(self, code_dim=120, n_class=1000, chn=96, debug=False):
super().__init__()
self.linear = nn.Linear(n_class, 128, bias=False)
if debug:
chn = 8
self.first_view = 16 * chn
self.G_linear = SpectralNorm(nn.Linear(20, 4 * 4 * 16 * chn))
z_dim = code_dim + 28
self.GBlock = nn.ModuleList([
GBlock(16 * chn, 16 * chn, n_class=n_class, z_dim=z_dim),
GBlock(16 * chn, 8 * chn, n_class=n_class, z_dim=z_dim),
GBlock(8 * chn, 4 * chn, n_class=n_class, z_dim=z_dim),
GBlock(4 * chn, 2 * chn, n_class=n_class, z_dim=z_dim),
GBlock(2 * chn, 1 * chn, n_class=n_class, z_dim=z_dim),
])
self.sa_id = 4
self.num_split = len(self.GBlock) + 1
self.attention = SelfAttention(2 * chn)
self.ScaledCrossReplicaBN = nn.BatchNorm2d(1 * chn, eps=1e-4)
self.colorize = SpectralNorm(nn.Conv2d(1 * chn, 3, [3, 3], padding=1))
def forward(self, input, class_id):
codes = torch.chunk(input, self.num_split, 1)
class_emb = self.linear(class_id) # 128
out = self.G_linear(codes[0])
out = out.view(-1, 4, 4, self.first_view).permute(0, 3, 1, 2)
for i, (code, GBlock) in enumerate(zip(codes[1:], self.GBlock)):
if i == self.sa_id:
out = self.attention(out)
condition = torch.cat([code, class_emb], 1)
out = GBlock(out, condition)
out = self.ScaledCrossReplicaBN(out)
out = F.relu(out)
out = self.colorize(out)
return torch.tanh(out)
class Generator256(nn.Module):
def __init__(self, code_dim=140, n_class=1000, chn=96, debug=False):
super().__init__()
self.linear = nn.Linear(n_class, 128, bias=False)
if debug:
chn = 8
self.first_view = 16 * chn
self.G_linear = SpectralNorm(nn.Linear(20, 4 * 4 * 16 * chn))
self.GBlock = nn.ModuleList([
GBlock(16 * chn, 16 * chn, n_class=n_class),
GBlock(16 * chn, 8 * chn, n_class=n_class),
GBlock(8 * chn, 8 * chn, n_class=n_class),
GBlock(8 * chn, 4 * chn, n_class=n_class),
GBlock(4 * chn, 2 * chn, n_class=n_class),
GBlock(2 * chn, 1 * chn, n_class=n_class),
])
self.sa_id = 5
self.num_split = len(self.GBlock) + 1
self.attention = SelfAttention(2 * chn)
self.ScaledCrossReplicaBN = nn.BatchNorm2d(1 * chn, eps=1e-4)
self.colorize = SpectralNorm(nn.Conv2d(1 * chn, 3, [3, 3], padding=1))
def forward(self, input, class_id):
codes = torch.chunk(input, self.num_split, 1)
class_emb = self.linear(class_id) # 128
out = self.G_linear(codes[0])
out = out.view(-1, 4, 4, self.first_view).permute(0, 3, 1, 2)
for i, (code, GBlock) in enumerate(zip(codes[1:], self.GBlock)):
if i == self.sa_id:
out = self.attention(out)
condition = torch.cat([code, class_emb], 1)
out = GBlock(out, condition)
out = self.ScaledCrossReplicaBN(out)
out = F.relu(out)
out = self.colorize(out)
return torch.tanh(out)
class Generator512(nn.Module):
def __init__(self, code_dim=128, n_class=1000, chn=96, debug=False):
super().__init__()
self.linear = nn.Linear(n_class, 128, bias=False)
if debug:
chn = 8
self.first_view = 16 * chn
self.G_linear = SpectralNorm(nn.Linear(16, 4 * 4 * 16 * chn))
z_dim = code_dim + 16
self.GBlock = nn.ModuleList([
GBlock(16 * chn, 16 * chn, n_class=n_class, z_dim=z_dim),
GBlock(16 * chn, 8 * chn, n_class=n_class, z_dim=z_dim),
GBlock(8 * chn, 8 * chn, n_class=n_class, z_dim=z_dim),
GBlock(8 * chn, 4 * chn, n_class=n_class, z_dim=z_dim),
GBlock(4 * chn, 2 * chn, n_class=n_class, z_dim=z_dim),
GBlock(2 * chn, 1 * chn, n_class=n_class, z_dim=z_dim),
GBlock(1 * chn, 1 * chn, n_class=n_class, z_dim=z_dim),
])
self.sa_id = 4
self.num_split = len(self.GBlock) + 1
self.attention = SelfAttention(4 * chn)
self.ScaledCrossReplicaBN = nn.BatchNorm2d(1 * chn)
self.colorize = SpectralNorm(nn.Conv2d(1 * chn, 3, [3, 3], padding=1))
def forward(self, input, class_id):
codes = torch.chunk(input, self.num_split, 1)
class_emb = self.linear(class_id) # 128
out = self.G_linear(codes[0])
out = out.view(-1, 4, 4, self.first_view).permute(0, 3, 1, 2)
for i, (code, GBlock) in enumerate(zip(codes[1:], self.GBlock)):
if i == self.sa_id:
out = self.attention(out)
condition = torch.cat([code, class_emb], 1)
out = GBlock(out, condition)
out = self.ScaledCrossReplicaBN(out)
out = F.relu(out)
out = self.colorize(out)
return torch.tanh(out)
class Discriminator(nn.Module):
def __init__(self, n_class=1000, chn=96, debug=False):
super().__init__()
def conv(in_channel, out_channel, downsample=True):
return GBlock(in_channel, out_channel, bn=False, upsample=False, downsample=downsample)
if debug:
chn = 8
self.debug = debug
self.pre_conv = nn.Sequential(
SpectralNorm(nn.Conv2d(3, 1 * chn, 3, padding=1)),
nn.ReLU(),
SpectralNorm(nn.Conv2d(1 * chn, 1 * chn, 3, padding=1)),
nn.AvgPool2d(2),
)
self.pre_skip = SpectralNorm(nn.Conv2d(3, 1 * chn, 1))
self.conv = nn.Sequential(
conv(1 * chn, 1 * chn, downsample=True),
conv(1 * chn, 2 * chn, downsample=True),
SelfAttention(2 * chn),
conv(2 * chn, 2 * chn, downsample=True),
conv(2 * chn, 4 * chn, downsample=True),
conv(4 * chn, 8 * chn, downsample=True),
conv(8 * chn, 8 * chn, downsample=True),
conv(8 * chn, 16 * chn, downsample=True),
conv(16 * chn, 16 * chn, downsample=False),
)
self.linear = SpectralNorm(nn.Linear(16 * chn, 1))
self.embed = nn.Embedding(n_class, 16 * chn)
self.embed.weight.data.uniform_(-0.1, 0.1)
self.embed = SpectralNorm(self.embed)
def forward(self, input, class_id):
out = self.pre_conv(input)
out += self.pre_skip(F.avg_pool2d(input, 2))
out = self.conv(out)
out = F.relu(out)
out = out.view(out.size(0), out.size(1), -1)
out = out.sum(2)
out_linear = self.linear(out).squeeze(1)
embed = self.embed(class_id)
prod = (out * embed).sum(1)
return out_linear + prod
================================================
FILE: BigGAN_utils/TFHub/converter.py
================================================
"""Utilities for converting TFHub BigGAN generator weights to PyTorch.
Recommended usage:
To convert all BigGAN variants and generate test samples, use:
```bash
CUDA_VISIBLE_DEVICES=0 python converter.py --generate_samples
```
See `parse_args` for additional options.
"""
import argparse
import os
import sys
import h5py
import torch
import torch.nn as nn
from torchvision.utils import save_image
import tensorflow as tf
import tensorflow_hub as hub
import parse
# import reference biggan from this folder
import biggan_v1 as biggan_for_conversion
# Import model from main folder
sys.path.append('..')
import BigGAN
DEVICE = 'cuda'
HDF5_TMPL = 'biggan-{}.h5'
PTH_TMPL = 'biggan-{}.pth'
MODULE_PATH_TMPL = 'https://tfhub.dev/deepmind/biggan-{}/2'
Z_DIMS = {
128: 120,
256: 140,
512: 128}
RESOLUTIONS = list(Z_DIMS)
def dump_tfhub_to_hdf5(module_path, hdf5_path, redownload=False):
"""Loads TFHub weights and saves them to intermediate HDF5 file.
Args:
module_path ([Path-like]): Path to TFHub module.
hdf5_path ([Path-like]): Path to output HDF5 file.
Returns:
[h5py.File]: Loaded hdf5 file containing module weights.
"""
if os.path.exists(hdf5_path) and (not redownload):
print('Loading BigGAN hdf5 file from:', hdf5_path)
return h5py.File(hdf5_path, 'r')
print('Loading BigGAN module from:', module_path)
tf.reset_default_graph()
hub.Module(module_path)
print('Loaded BigGAN module from:', module_path)
initializer = tf.global_variables_initializer()
sess = tf.Session()
sess.run(initializer)
print('Saving BigGAN weights to :', hdf5_path)
h5f = h5py.File(hdf5_path, 'w')
for var in tf.global_variables():
val = sess.run(var)
h5f.create_dataset(var.name, data=val)
print(f'Saving {var.name} with shape {val.shape}')
h5f.close()
return h5py.File(hdf5_path, 'r')
class TFHub2Pytorch(object):
TF_ROOT = 'module'
NUM_GBLOCK = {
128: 5,
256: 6,
512: 7
}
w = 'w'
b = 'b'
u = 'u0'
v = 'u1'
gamma = 'gamma'
beta = 'beta'
def __init__(self, state_dict, tf_weights, resolution=256, load_ema=True, verbose=False):
self.state_dict = state_dict
self.tf_weights = tf_weights
self.resolution = resolution
self.verbose = verbose
if load_ema:
for name in ['w', 'b', 'gamma', 'beta']:
setattr(self, name, getattr(self, name) + '/ema_b999900')
def load(self):
self.load_generator()
return self.state_dict
def load_generator(self):
GENERATOR_ROOT = os.path.join(self.TF_ROOT, 'Generator')
for i in range(self.NUM_GBLOCK[self.resolution]):
name_tf = os.path.join(GENERATOR_ROOT, 'GBlock')
name_tf += f'_{i}' if i != 0 else ''
self.load_GBlock(f'GBlock.{i}.', name_tf)
self.load_attention('attention.', os.path.join(GENERATOR_ROOT, 'attention'))
self.load_linear('linear', os.path.join(self.TF_ROOT, 'linear'), bias=False)
self.load_snlinear('G_linear', os.path.join(GENERATOR_ROOT, 'G_Z', 'G_linear'))
self.load_colorize('colorize', os.path.join(GENERATOR_ROOT, 'conv_2d'))
self.load_ScaledCrossReplicaBNs('ScaledCrossReplicaBN',
os.path.join(GENERATOR_ROOT, 'ScaledCrossReplicaBN'))
def load_linear(self, name_pth, name_tf, bias=True):
self.state_dict[name_pth + '.weight'] = self.load_tf_tensor(name_tf, self.w).permute(1, 0)
if bias:
self.state_dict[name_pth + '.bias'] = self.load_tf_tensor(name_tf, self.b)
def load_snlinear(self, name_pth, name_tf, bias=True):
self.state_dict[name_pth + '.module.weight_u'] = self.load_tf_tensor(name_tf, self.u).squeeze()
self.state_dict[name_pth + '.module.weight_v'] = self.load_tf_tensor(name_tf, self.v).squeeze()
self.state_dict[name_pth + '.module.weight_bar'] = self.load_tf_tensor(name_tf, self.w).permute(1, 0)
if bias:
self.state_dict[name_pth + '.module.bias'] = self.load_tf_tensor(name_tf, self.b)
def load_colorize(self, name_pth, name_tf):
self.load_snconv(name_pth, name_tf)
def load_GBlock(self, name_pth, name_tf):
self.load_convs(name_pth, name_tf)
self.load_HyperBNs(name_pth, name_tf)
def load_convs(self, name_pth, name_tf):
self.load_snconv(name_pth + 'conv0', os.path.join(name_tf, 'conv0'))
self.load_snconv(name_pth + 'conv1', os.path.join(name_tf, 'conv1'))
self.load_snconv(name_pth + 'conv_sc', os.path.join(name_tf, 'conv_sc'))
def load_snconv(self, name_pth, name_tf, bias=True):
if self.verbose:
print(f'loading: {name_pth} from {name_tf}')
self.state_dict[name_pth + '.module.weight_u'] = self.load_tf_tensor(name_tf, self.u).squeeze()
self.state_dict[name_pth + '.module.weight_v'] = self.load_tf_tensor(name_tf, self.v).squeeze()
self.state_dict[name_pth + '.module.weight_bar'] = self.load_tf_tensor(name_tf, self.w).permute(3, 2, 0, 1)
if bias:
self.state_dict[name_pth + '.module.bias'] = self.load_tf_tensor(name_tf, self.b).squeeze()
def load_conv(self, name_pth, name_tf, bias=True):
self.state_dict[name_pth + '.weight_u'] = self.load_tf_tensor(name_tf, self.u).squeeze()
self.state_dict[name_pth + '.weight_v'] = self.load_tf_tensor(name_tf, self.v).squeeze()
self.state_dict[name_pth + '.weight_bar'] = self.load_tf_tensor(name_tf, self.w).permute(3, 2, 0, 1)
if bias:
self.state_dict[name_pth + '.bias'] = self.load_tf_tensor(name_tf, self.b)
def load_HyperBNs(self, name_pth, name_tf):
self.load_HyperBN(name_pth + 'HyperBN', os.path.join(name_tf, 'HyperBN'))
self.load_HyperBN(name_pth + 'HyperBN_1', os.path.join(name_tf, 'HyperBN_1'))
def load_ScaledCrossReplicaBNs(self, name_pth, name_tf):
self.state_dict[name_pth + '.bias'] = self.load_tf_tensor(name_tf, self.beta).squeeze()
self.state_dict[name_pth + '.weight'] = self.load_tf_tensor(name_tf, self.gamma).squeeze()
self.state_dict[name_pth + '.running_mean'] = self.load_tf_tensor(name_tf + 'bn', 'accumulated_mean')
self.state_dict[name_pth + '.running_var'] = self.load_tf_tensor(name_tf + 'bn', 'accumulated_var')
self.state_dict[name_pth + '.num_batches_tracked'] = torch.tensor(
self.tf_weights[os.path.join(name_tf + 'bn', 'accumulation_counter:0')][()], dtype=torch.float32)
def load_HyperBN(self, name_pth, name_tf):
if self.verbose:
print(f'loading: {name_pth} from {name_tf}')
beta = name_pth + '.beta_embed.module'
gamma = name_pth + '.gamma_embed.module'
self.state_dict[beta + '.weight_u'] = self.load_tf_tensor(os.path.join(name_tf, 'beta'), self.u).squeeze()
self.state_dict[gamma + '.weight_u'] = self.load_tf_tensor(os.path.join(name_tf, 'gamma'), self.u).squeeze()
self.state_dict[beta + '.weight_v'] = self.load_tf_tensor(os.path.join(name_tf, 'beta'), self.v).squeeze()
self.state_dict[gamma + '.weight_v'] = self.load_tf_tensor(os.path.join(name_tf, 'gamma'), self.v).squeeze()
self.state_dict[beta + '.weight_bar'] = self.load_tf_tensor(os.path.join(name_tf, 'beta'), self.w).permute(1, 0)
self.state_dict[gamma +
'.weight_bar'] = self.load_tf_tensor(os.path.join(name_tf, 'gamma'), self.w).permute(1, 0)
cr_bn_name = name_tf.replace('HyperBN', 'CrossReplicaBN')
self.state_dict[name_pth + '.bn.running_mean'] = self.load_tf_tensor(cr_bn_name, 'accumulated_mean')
self.state_dict[name_pth + '.bn.running_var'] = self.load_tf_tensor(cr_bn_name, 'accumulated_var')
self.state_dict[name_pth + '.bn.num_batches_tracked'] = torch.tensor(
self.tf_weights[os.path.join(cr_bn_name, 'accumulation_counter:0')][()], dtype=torch.float32)
def load_attention(self, name_pth, name_tf):
self.load_snconv(name_pth + 'theta', os.path.join(name_tf, 'theta'), bias=False)
self.load_snconv(name_pth + 'phi', os.path.join(name_tf, 'phi'), bias=False)
self.load_snconv(name_pth + 'g', os.path.join(name_tf, 'g'), bias=False)
self.load_snconv(name_pth + 'o_conv', os.path.join(name_tf, 'o_conv'), bias=False)
self.state_dict[name_pth + 'gamma'] = self.load_tf_tensor(name_tf, self.gamma)
def load_tf_tensor(self, prefix, var, device='0'):
name = os.path.join(prefix, var) + f':{device}'
return torch.from_numpy(self.tf_weights[name][:])
# Convert from v1: This function maps
def convert_from_v1(hub_dict, resolution=128):
weightname_dict = {'weight_u': 'u0', 'weight_bar': 'weight', 'bias': 'bias'}
convnum_dict = {'conv0': 'conv1', 'conv1': 'conv2', 'conv_sc': 'conv_sc'}
attention_blocknum = {128: 3, 256: 4, 512: 3}[resolution]
hub2me = {'linear.weight': 'shared.weight', # This is actually the shared weight
# Linear stuff
'G_linear.module.weight_bar': 'linear.weight',
'G_linear.module.bias': 'linear.bias',
'G_linear.module.weight_u': 'linear.u0',
# output layer stuff
'ScaledCrossReplicaBN.weight': 'output_layer.0.gain',
'ScaledCrossReplicaBN.bias': 'output_layer.0.bias',
'ScaledCrossReplicaBN.running_mean': 'output_layer.0.stored_mean',
'ScaledCrossReplicaBN.running_var': 'output_layer.0.stored_var',
'colorize.module.weight_bar': 'output_layer.2.weight',
'colorize.module.bias': 'output_layer.2.bias',
'colorize.module.weight_u': 'output_layer.2.u0',
# Attention stuff
'attention.gamma': 'blocks.%d.1.gamma' % attention_blocknum,
'attention.theta.module.weight_u': 'blocks.%d.1.theta.u0' % attention_blocknum,
'attention.theta.module.weight_bar': 'blocks.%d.1.theta.weight' % attention_blocknum,
'attention.phi.module.weight_u': 'blocks.%d.1.phi.u0' % attention_blocknum,
'attention.phi.module.weight_bar': 'blocks.%d.1.phi.weight' % attention_blocknum,
'attention.g.module.weight_u': 'blocks.%d.1.g.u0' % attention_blocknum,
'attention.g.module.weight_bar': 'blocks.%d.1.g.weight' % attention_blocknum,
'attention.o_conv.module.weight_u': 'blocks.%d.1.o.u0' % attention_blocknum,
'attention.o_conv.module.weight_bar':'blocks.%d.1.o.weight' % attention_blocknum,
}
# Loop over the hub dict and build the hub2me map
for name in hub_dict.keys():
if 'GBlock' in name:
if 'HyperBN' not in name: # it's a conv
out = parse.parse('GBlock.{:d}.{}.module.{}',name)
blocknum, convnum, weightname = out
if weightname not in weightname_dict:
continue # else hyperBN in
out_name = 'blocks.%d.0.%s.%s' % (blocknum, convnum_dict[convnum], weightname_dict[weightname]) # Increment conv number by 1
else: # hyperbn not conv
BNnum = 2 if 'HyperBN_1' in name else 1
if 'embed' in name:
out = parse.parse('GBlock.{:d}.{}.module.{}',name)
blocknum, gamma_or_beta, weightname = out
if weightname not in weightname_dict: # Ignore weight_v
continue
out_name = 'blocks.%d.0.bn%d.%s.%s' % (blocknum, BNnum, 'gain' if 'gamma' in gamma_or_beta else 'bias', weightname_dict[weightname])
else:
out = parse.parse('GBlock.{:d}.{}.bn.{}',name)
blocknum, dummy, mean_or_var = out
if 'num_batches_tracked' in mean_or_var:
continue
out_name = 'blocks.%d.0.bn%d.%s' % (blocknum, BNnum, 'stored_mean' if 'mean' in mean_or_var else 'stored_var')
hub2me[name] = out_name
# Invert the hub2me map
me2hub = {hub2me[item]: item for item in hub2me}
new_dict = {}
dimz_dict = {128: 20, 256: 20, 512:16}
for item in me2hub:
# Swap input dim ordering on batchnorm bois to account for my arbitrary change of ordering when concatenating Ys and Zs
if ('bn' in item and 'weight' in item) and ('gain' in item or 'bias' in item) and ('output_layer' not in item):
new_dict[item] = torch.cat([hub_dict[me2hub[item]][:, -128:], hub_dict[me2hub[item]][:, :dimz_dict[resolution]]], 1)
# Reshape the first linear weight, bias, and u0
elif item == 'linear.weight':
new_dict[item] = hub_dict[me2hub[item]].contiguous().view(4, 4, 96 * 16, -1).permute(2,0,1,3).contiguous().view(-1,dimz_dict[resolution])
elif item == 'linear.bias':
new_dict[item] = hub_dict[me2hub[item]].view(4, 4, 96 * 16).permute(2,0,1).contiguous().view(-1)
elif item == 'linear.u0':
new_dict[item] = hub_dict[me2hub[item]].view(4, 4, 96 * 16).permute(2,0,1).contiguous().view(1, -1)
elif me2hub[item] == 'linear.weight': # THIS IS THE SHARED WEIGHT NOT THE FIRST LINEAR LAYER
# Transpose shared weight so that it's an embedding
new_dict[item] = hub_dict[me2hub[item]].t()
elif 'weight_u' in me2hub[item]: # Unsqueeze u0s
new_dict[item] = hub_dict[me2hub[item]].unsqueeze(0)
else:
new_dict[item] = hub_dict[me2hub[item]]
return new_dict
def get_config(resolution):
attn_dict = {128: '64', 256: '128', 512: '64'}
dim_z_dict = {128: 120, 256: 140, 512: 128}
config = {'G_param': 'SN', 'D_param': 'SN',
'G_ch': 96, 'D_ch': 96,
'D_wide': True, 'G_shared': True,
'shared_dim': 128, 'dim_z': dim_z_dict[resolution],
'hier': True, 'cross_replica': False,
'mybn': False, 'G_activation': nn.ReLU(inplace=True),
'G_attn': attn_dict[resolution],
'norm_style': 'bn',
'G_init': 'ortho', 'skip_init': True, 'no_optim': True,
'G_fp16': False, 'G_mixed_precision': False,
'accumulate_stats': False, 'num_standing_accumulations': 16,
'G_eval_mode': True,
'BN_eps': 1e-04, 'SN_eps': 1e-04,
'num_G_SVs': 1, 'num_G_SV_itrs': 1, 'resolution': resolution,
'n_classes': 1000}
return config
def convert_biggan(resolution, weight_dir, redownload=False, no_ema=False, verbose=False):
module_path = MODULE_PATH_TMPL.format(resolution)
hdf5_path = os.path.join(weight_dir, HDF5_TMPL.format(resolution))
pth_path = os.path.join(weight_dir, PTH_TMPL.format(resolution))
tf_weights = dump_tfhub_to_hdf5(module_path, hdf5_path, redownload=redownload)
G_temp = getattr(biggan_for_conversion, f'Generator{resolution}')()
state_dict_temp = G_temp.state_dict()
converter = TFHub2Pytorch(state_dict_temp, tf_weights, resolution=resolution,
load_ema=(not no_ema), verbose=verbose)
state_dict_v1 = converter.load()
state_dict = convert_from_v1(state_dict_v1, resolution)
# Get the config, build the model
config = get_config(resolution)
G = BigGAN.Generator(**config)
G.load_state_dict(state_dict, strict=False) # Ignore missing sv0 entries
torch.save(state_dict, pth_path)
# output_location ='pretrained_weights/TFHub-PyTorch-128.pth'
return G
def generate_sample(G, z_dim, batch_size, filename, parallel=False):
G.eval()
G.to(DEVICE)
with torch.no_grad():
z = torch.randn(batch_size, G.dim_z).to(DEVICE)
y = torch.randint(low=0, high=1000, size=(batch_size,),
device=DEVICE, dtype=torch.int64, requires_grad=False)
if parallel:
images = nn.parallel.data_parallel(G, (z, G.shared(y)))
else:
images = G(z, G.shared(y))
save_image(images, filename, scale_each=True, normalize=True)
def parse_args():
usage = 'Parser for conversion script.'
parser = argparse.ArgumentParser(description=usage)
parser.add_argument(
'--resolution', '-r', type=int, default=None, choices=[128, 256, 512],
help='Resolution of TFHub module to convert. Converts all resolutions if None.')
parser.add_argument(
'--redownload', action='store_true', default=False,
help='Redownload weights and overwrite current hdf5 file, if present.')
parser.add_argument(
'--weights_dir', type=str, default='pretrained_weights')
parser.add_argument(
'--samples_dir', type=str, default='pretrained_samples')
parser.add_argument(
'--no_ema', action='store_true', default=False,
help='Do not load ema weights.')
parser.add_argument(
'--verbose', action='store_true', default=False,
help='Additionally logging.')
parser.add_argument(
'--generate_samples', action='store_true', default=False,
help='Generate test sample with pretrained model.')
parser.add_argument(
'--batch_size', type=int, default=64,
help='Batch size used for test sample.')
parser.add_argument(
'--parallel', action='store_true', default=False,
help='Parallelize G?')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
os.makedirs(args.weights_dir, exist_ok=True)
os.makedirs(args.samples_dir, exist_ok=True)
if args.resolution is not None:
G = convert_biggan(args.resolution, args.weights_dir,
redownload=args.redownload,
no_ema=args.no_ema, verbose=args.verbose)
if args.generate_samples:
filename = os.path.join(args.samples_dir, f'biggan{args.resolution}_samples.jpg')
print('Generating samples...')
generate_sample(G, Z_DIMS[args.resolution], args.batch_size, filename, args.parallel)
else:
for res in RESOLUTIONS:
G = convert_biggan(res, args.weights_dir,
redownload=args.redownload,
no_ema=args.no_ema, verbose=args.verbose)
if args.generate_samples:
filename = os.path.join(args.samples_dir, f'biggan{res}_samples.jpg')
print('Generating samples...')
generate_sample(G, Z_DIMS[res], args.batch_size, filename, args.parallel)
================================================
FILE: BigGAN_utils/__init__.py
================================================
import sys
sys.path.append('./BigGAN_utils/')
================================================
FILE: BigGAN_utils/animal_hash.py
================================================
c = ['Aardvark', 'Abyssinian', 'Affenpinscher', 'Akbash', 'Akita', 'Albatross',
'Alligator', 'Alpaca', 'Angelfish', 'Ant', 'Anteater', 'Antelope', 'Ape',
'Armadillo', 'Ass', 'Avocet', 'Axolotl', 'Baboon', 'Badger', 'Balinese',
'Bandicoot', 'Barb', 'Barnacle', 'Barracuda', 'Bat', 'Beagle', 'Bear',
'Beaver', 'Bee', 'Beetle', 'Binturong', 'Bird', 'Birman', 'Bison',
'Bloodhound', 'Boar', 'Bobcat', 'Bombay', 'Bongo', 'Bonobo', 'Booby',
'Budgerigar', 'Buffalo', 'Bulldog', 'Bullfrog', 'Burmese', 'Butterfly',
'Caiman', 'Camel', 'Capybara', 'Caracal', 'Caribou', 'Cassowary', 'Cat',
'Caterpillar', 'Catfish', 'Cattle', 'Centipede', 'Chameleon', 'Chamois',
'Cheetah', 'Chicken', 'Chihuahua', 'Chimpanzee', 'Chinchilla', 'Chinook',
'Chipmunk', 'Chough', 'Cichlid', 'Clam', 'Coati', 'Cobra', 'Cockroach',
'Cod', 'Collie', 'Coral', 'Cormorant', 'Cougar', 'Cow', 'Coyote',
'Crab', 'Crane', 'Crocodile', 'Crow', 'Curlew', 'Cuscus', 'Cuttlefish',
'Dachshund', 'Dalmatian', 'Deer', 'Dhole', 'Dingo', 'Dinosaur', 'Discus',
'Dodo', 'Dog', 'Dogball', 'Dogfish', 'Dolphin', 'Donkey', 'Dormouse',
'Dove', 'Dragonfly', 'Drever', 'Duck', 'Dugong', 'Dunker', 'Dunlin',
'Eagle', 'Earwig', 'Echidna', 'Eel', 'Eland', 'Elephant', 'ElephantSeal',
'Elk', 'Emu', 'Falcon', 'Ferret', 'Finch', 'Fish', 'Flamingo', 'Flounder',
'Fly', 'Fossa', 'Fox', 'Frigatebird', 'Frog', 'Galago', 'Gar', 'Gaur',
'Gazelle', 'Gecko', 'Gerbil', 'Gharial', 'GiantPanda', 'Gibbon', 'Giraffe',
'Gnat', 'Gnu', 'Goat', 'Goldfinch', 'Goldfish', 'Goose', 'Gopher',
'Gorilla', 'Goshawk', 'Grasshopper', 'Greyhound', 'Grouse', 'Guanaco',
'GuineaFowl', 'GuineaPig', 'Gull', 'Guppy', 'Hamster', 'Hare', 'Harrier',
'Havanese', 'Hawk', 'Hedgehog', 'Heron', 'Herring', 'Himalayan',
'Hippopotamus', 'Hornet', 'Horse', 'Human', 'Hummingbird', 'Hyena',
'Ibis', 'Iguana', 'Impala', 'Indri', 'Insect', 'Jackal', 'Jaguar',
'Javanese', 'Jay', 'Jellyfish', 'Kakapo', 'Kangaroo', 'Kingfisher',
'Kiwi', 'Koala', 'KomodoDragon', 'Kouprey', 'Kudu', 'Labradoodle',
'Ladybird', 'Lapwing', 'Lark', 'Lemming', 'Lemur', 'Leopard', 'Liger',
'Lion', 'Lionfish', 'Lizard', 'Llama', 'Lobster', 'Locust', 'Loris',
'Louse', 'Lynx', 'Lyrebird', 'Macaw', 'Magpie', 'Mallard', 'Maltese',
'Manatee', 'Mandrill', 'Markhor', 'Marten', 'Mastiff', 'Mayfly', 'Meerkat',
'Millipede', 'Mink', 'Mole', 'Molly', 'Mongoose', 'Mongrel', 'Monkey',
'Moorhen', 'Moose', 'Mosquito', 'Moth', 'Mouse', 'Mule', 'Narwhal',
'Neanderthal', 'Newfoundland', 'Newt', 'Nightingale', 'Numbat', 'Ocelot',
'Octopus', 'Okapi', 'Olm', 'Opossum', 'Orang-utan', 'Oryx', 'Ostrich',
'Otter', 'Owl', 'Ox', 'Oyster', 'Pademelon', 'Panther', 'Parrot',
'Partridge', 'Peacock', 'Peafowl', 'Pekingese', 'Pelican', 'Penguin',
'Persian', 'Pheasant', 'Pig', 'Pigeon', 'Pika', 'Pike', 'Piranha',
'Platypus', 'Pointer', 'Pony', 'Poodle', 'Porcupine', 'Porpoise',
'Possum', 'PrairieDog', 'Prawn', 'Puffin', 'Pug', 'Puma', 'Quail',
'Quelea', 'Quetzal', 'Quokka', 'Quoll', 'Rabbit', 'Raccoon', 'Ragdoll',
'Rail', 'Ram', 'Rat', 'Rattlesnake', 'Raven', 'RedDeer', 'RedPanda',
'Reindeer', 'Rhinoceros', 'Robin', 'Rook', 'Rottweiler', 'Ruff',
'Salamander', 'Salmon', 'SandDollar', 'Sandpiper', 'Saola',
'Sardine', 'Scorpion', 'SeaLion', 'SeaUrchin', 'Seahorse',
'Seal', 'Serval', 'Shark', 'Sheep', 'Shrew', 'Shrimp', 'Siamese',
'Siberian', 'Skunk', 'Sloth', 'Snail', 'Snake', 'Snowshoe', 'Somali',
'Sparrow', 'Spider', 'Sponge', 'Squid', 'Squirrel', 'Starfish', 'Starling',
'Stingray', 'Stinkbug', 'Stoat', 'Stork', 'Swallow', 'Swan', 'Tang',
'Tapir', 'Tarsier', 'Termite', 'Tetra', 'Tiffany', 'Tiger', 'Toad',
'Tortoise', 'Toucan', 'Tropicbird', 'Trout', 'Tuatara', 'Turkey',
'Turtle', 'Uakari', 'Uguisu', 'Umbrellabird', 'Viper', 'Vulture',
'Wallaby', 'Walrus', 'Warthog', 'Wasp', 'WaterBuffalo', 'Weasel',
'Whale', 'Whippet', 'Wildebeest', 'Wolf', 'Wolverine', 'Wombat',
'Woodcock', 'Woodlouse', 'Woodpecker', 'Worm', 'Wrasse', 'Wren',
'Yak', 'Zebra', 'Zebu', 'Zonkey']
a = ['able', 'above', 'absent', 'absolute', 'abstract', 'abundant', 'academic',
'acceptable', 'accepted', 'accessible', 'accurate', 'accused', 'active',
'actual', 'acute', 'added', 'additional', 'adequate', 'adjacent',
'administrative', 'adorable', 'advanced', 'adverse', 'advisory',
'aesthetic', 'afraid', 'african', 'aggregate', 'aggressive', 'agreeable',
'agreed', 'agricultural', 'alert', 'alive', 'alleged', 'allied', 'alone',
'alright', 'alternative', 'amateur', 'amazing', 'ambitious', 'american',
'amused', 'ancient', 'angry', 'annoyed', 'annual', 'anonymous', 'anxious',
'appalling', 'apparent', 'applicable', 'appropriate', 'arab', 'arbitrary',
'architectural', 'armed', 'arrogant', 'artificial', 'artistic', 'ashamed',
'asian', 'asleep', 'assistant', 'associated', 'atomic', 'attractive',
'australian', 'automatic', 'autonomous', 'available', 'average',
'awake', 'aware', 'awful', 'awkward', 'back', 'bad', 'balanced', 'bare',
'basic', 'beautiful', 'beneficial', 'better', 'bewildered', 'big',
'binding', 'biological', 'bitter', 'bizarre', 'black', 'blank', 'blind',
'blonde', 'bloody', 'blue', 'blushing', 'boiling', 'bold', 'bored',
'boring', 'bottom', 'brainy', 'brave', 'breakable', 'breezy', 'brief',
'bright', 'brilliant', 'british', 'broad', 'broken', 'brown', 'bumpy',
'burning', 'busy', 'calm', 'canadian', 'capable', 'capitalist', 'careful',
'casual', 'catholic', 'causal', 'cautious', 'central', 'certain',
'changing', 'characteristic', 'charming', 'cheap', 'cheerful', 'chemical',
'chief', 'chilly', 'chinese', 'chosen', 'christian', 'chronic', 'chubby',
'circular', 'civic', 'civil', 'civilian', 'classic', 'classical', 'clean',
'clear', 'clever', 'clinical', 'close', 'closed', 'cloudy', 'clumsy',
'coastal', 'cognitive', 'coherent', 'cold', 'collective', 'colonial',
'colorful', 'colossal', 'coloured', 'colourful', 'combative', 'combined',
'comfortable', 'coming', 'commercial', 'common', 'communist', 'compact',
'comparable', 'comparative', 'compatible', 'competent', 'competitive',
'complete', 'complex', 'complicated', 'comprehensive', 'compulsory',
'conceptual', 'concerned', 'concrete', 'condemned', 'confident',
'confidential', 'confused', 'conscious', 'conservation', 'conservative',
'considerable', 'consistent', 'constant', 'constitutional',
'contemporary', 'content', 'continental', 'continued', 'continuing',
'continuous', 'controlled', 'controversial', 'convenient', 'conventional',
'convinced', 'convincing', 'cooing', 'cool', 'cooperative', 'corporate',
'correct', 'corresponding', 'costly', 'courageous', 'crazy', 'creative',
'creepy', 'criminal', 'critical', 'crooked', 'crowded', 'crucial',
'crude', 'cruel', 'cuddly', 'cultural', 'curious', 'curly', 'current',
'curved', 'cute', 'daily', 'damaged', 'damp', 'dangerous', 'dark', 'dead',
'deaf', 'deafening', 'dear', 'decent', 'decisive', 'deep', 'defeated',
'defensive', 'defiant', 'definite', 'deliberate', 'delicate', 'delicious',
'delighted', 'delightful', 'democratic', 'dependent', 'depressed',
'desirable', 'desperate', 'detailed', 'determined', 'developed',
'developing', 'devoted', 'different', 'difficult', 'digital', 'diplomatic',
'direct', 'dirty', 'disabled', 'disappointed', 'disastrous',
'disciplinary', 'disgusted', 'distant', 'distinct', 'distinctive',
'distinguished', 'disturbed', 'disturbing', 'diverse', 'divine', 'dizzy',
'domestic', 'dominant', 'double', 'doubtful', 'drab', 'dramatic',
'dreadful', 'driving', 'drunk', 'dry', 'dual', 'due', 'dull', 'dusty',
'dutch', 'dying', 'dynamic', 'eager', 'early', 'eastern', 'easy',
'economic', 'educational', 'eerie', 'effective', 'efficient',
'elaborate', 'elated', 'elderly', 'eldest', 'electoral', 'electric',
'electrical', 'electronic', 'elegant', 'eligible', 'embarrassed',
'embarrassing', 'emotional', 'empirical', 'empty', 'enchanting',
'encouraging', 'endless', 'energetic', 'english', 'enormous',
'enthusiastic', 'entire', 'entitled', 'envious', 'environmental', 'equal',
'equivalent', 'essential', 'established', 'estimated', 'ethical',
'ethnic', 'european', 'eventual', 'everyday', 'evident', 'evil',
'evolutionary', 'exact', 'excellent', 'exceptional', 'excess',
'excessive', 'excited', 'exciting', 'exclusive', 'existing', 'exotic',
'expected', 'expensive', 'experienced', 'experimental', 'explicit',
'extended', 'extensive', 'external', 'extra', 'extraordinary', 'extreme',
'exuberant', 'faint', 'fair', 'faithful', 'familiar', 'famous', 'fancy',
'fantastic', 'far', 'fascinating', 'fashionable', 'fast', 'fat', 'fatal',
'favourable', 'favourite', 'federal', 'fellow', 'female', 'feminist',
'few', 'fierce', 'filthy', 'final', 'financial', 'fine', 'firm', 'fiscal',
'fit', 'fixed', 'flaky', 'flat', 'flexible', 'fluffy', 'fluttering',
'flying', 'following', 'fond', 'foolish', 'foreign', 'formal',
'formidable', 'forthcoming', 'fortunate', 'forward', 'fragile',
'frail', 'frantic', 'free', 'french', 'frequent', 'fresh', 'friendly',
'frightened', 'front', 'frozen', 'fucking', 'full', 'full-time', 'fun',
'functional', 'fundamental', 'funny', 'furious', 'future', 'fuzzy',
'gastric', 'gay', 'general', 'generous', 'genetic', 'gentle', 'genuine',
'geographical', 'german', 'giant', 'gigantic', 'given', 'glad',
'glamorous', 'gleaming', 'global', 'glorious', 'golden', 'good',
'gorgeous', 'gothic', 'governing', 'graceful', 'gradual', 'grand',
'grateful', 'greasy', 'great', 'greek', 'green', 'grey', 'grieving',
'grim', 'gross', 'grotesque', 'growing', 'grubby', 'grumpy', 'guilty',
'handicapped', 'handsome', 'happy', 'hard', 'harsh', 'head', 'healthy',
'heavy', 'helpful', 'helpless', 'hidden', 'high', 'high-pitched',
'hilarious', 'hissing', 'historic', 'historical', 'hollow', 'holy',
'homeless', 'homely', 'hon', 'honest', 'horizontal', 'horrible',
'hostile', 'hot', 'huge', 'human', 'hungry', 'hurt', 'hushed', 'husky',
'icy', 'ideal', 'identical', 'ideological', 'ill', 'illegal',
'imaginative', 'immediate', 'immense', 'imperial', 'implicit',
'important', 'impossible', 'impressed', 'impressive', 'improved',
'inadequate', 'inappropriate', 'inc', 'inclined', 'increased',
'increasing', 'incredible', 'independent', 'indian', 'indirect',
'individual', 'industrial', 'inevitable', 'influential', 'informal',
'inherent', 'initial', 'injured', 'inland', 'inner', 'innocent',
'innovative', 'inquisitive', 'instant', 'institutional', 'insufficient',
'intact', 'integral', 'integrated', 'intellectual', 'intelligent',
'intense', 'intensive', 'interested', 'interesting', 'interim',
'interior', 'intermediate', 'internal', 'international', 'intimate',
'invisible', 'involved', 'iraqi', 'irish', 'irrelevant', 'islamic',
'isolated', 'israeli', 'italian', 'itchy', 'japanese', 'jealous',
'jewish', 'jittery', 'joint', 'jolly', 'joyous', 'judicial', 'juicy',
'junior', 'just', 'keen', 'key', 'kind', 'known', 'korean', 'labour',
'large', 'large-scale', 'late', 'latin', 'lazy', 'leading', 'left',
'legal', 'legislative', 'legitimate', 'lengthy', 'lesser', 'level',
'lexical', 'liable', 'liberal', 'light', 'like', 'likely', 'limited',
'linear', 'linguistic', 'liquid', 'literary', 'little', 'live', 'lively',
'living', 'local', 'logical', 'lonely', 'long', 'long-term', 'loose',
'lost', 'loud', 'lovely', 'low', 'loyal', 'ltd', 'lucky', 'mad',
'magenta', 'magic', 'magnetic', 'magnificent', 'main', 'major', 'male',
'mammoth', 'managerial', 'managing', 'manual', 'many', 'marginal',
'marine', 'marked', 'married', 'marvellous', 'marxist', 'mass', 'massive',
'mathematical', 'mature', 'maximum', 'mean', 'meaningful', 'mechanical',
'medical', 'medieval', 'melodic', 'melted', 'mental', 'mere',
'metropolitan', 'mid', 'middle', 'middle-class', 'mighty', 'mild',
'military', 'miniature', 'minimal', 'minimum', 'ministerial', 'minor',
'miserable', 'misleading', 'missing', 'misty', 'mixed', 'moaning',
'mobile', 'moderate', 'modern', 'modest', 'molecular', 'monetary',
'monthly', 'moral', 'motionless', 'muddy', 'multiple', 'mushy',
'musical', 'mute', 'mutual', 'mysterious', 'naked', 'narrow', 'nasty',
'national', 'native', 'natural', 'naughty', 'naval', 'near', 'nearby',
'neat', 'necessary', 'negative', 'neighbouring', 'nervous', 'net',
'neutral', 'new', 'nice', 'nineteenth-century', 'noble', 'noisy',
'normal', 'northern', 'nosy', 'notable', 'novel', 'nuclear', 'numerous',
'nursing', 'nutritious', 'nutty', 'obedient', 'objective', 'obliged',
'obnoxious', 'obvious', 'occasional', 'occupational', 'odd', 'official',
'ok', 'okay', 'old', 'old-fashioned', 'olympic', 'only', 'open',
'operational', 'opposite', 'optimistic', 'oral', 'orange', 'ordinary',
'organic', 'organisational', 'original', 'orthodox', 'other', 'outdoor',
'outer', 'outrageous', 'outside', 'outstanding', 'overall', 'overseas',
'overwhelming', 'painful', 'pale', 'palestinian', 'panicky', 'parallel',
'parental', 'parliamentary', 'part-time', 'partial', 'particular',
'passing', 'passive', 'past', 'patient', 'payable', 'peaceful',
'peculiar', 'perfect', 'permanent', 'persistent', 'personal', 'petite',
'philosophical', 'physical', 'pink', 'plain', 'planned', 'plastic',
'pleasant', 'pleased', 'poised', 'polish', 'polite', 'political', 'poor',
'popular', 'positive', 'possible', 'post-war', 'potential', 'powerful',
'practical', 'precious', 'precise', 'preferred', 'pregnant',
'preliminary', 'premier', 'prepared', 'present', 'presidential',
'pretty', 'previous', 'prickly', 'primary', 'prime', 'primitive',
'principal', 'printed', 'prior', 'private', 'probable', 'productive',
'professional', 'profitable', 'profound', 'progressive', 'prominent',
'promising', 'proper', 'proposed', 'prospective', 'protective',
'protestant', 'proud', 'provincial', 'psychiatric', 'psychological',
'public', 'puny', 'pure', 'purple', 'purring', 'puzzled', 'quaint',
'qualified', 'quick', 'quickest', 'quiet', 'racial', 'radical', 'rainy',
'random', 'rapid', 'rare', 'raspy', 'rational', 'ratty', 'raw', 'ready',
'real', 'realistic', 'rear', 'reasonable', 'recent', 'red', 'reduced',
'redundant', 'regional', 'registered', 'regular', 'regulatory', 'related',
'relative', 'relaxed', 'relevant', 'reliable', 'relieved', 'religious',
'reluctant', 'remaining', 'remarkable', 'remote', 'renewed',
'representative', 'repulsive', 'required', 'resident', 'residential',
'resonant', 'respectable', 'respective', 'responsible', 'resulting',
'retail', 'retired', 'revolutionary', 'rich', 'ridiculous', 'right',
'rigid', 'ripe', 'rising', 'rival', 'roasted', 'robust', 'rolling',
'roman', 'romantic', 'rotten', 'rough', 'round', 'royal', 'rubber',
'rude', 'ruling', 'running', 'rural', 'russian', 'sacred', 'sad', 'safe',
'salty', 'satisfactory', 'satisfied', 'scared', 'scary', 'scattered',
'scientific', 'scornful', 'scottish', 'scrawny', 'screeching',
'secondary', 'secret', 'secure', 'select', 'selected', 'selective',
'selfish', 'semantic', 'senior', 'sensible', 'sensitive', 'separate',
'serious', 'severe', 'sexual', 'shaggy', 'shaky', 'shallow', 'shared',
'sharp', 'sheer', 'shiny', 'shivering', 'shocked', 'short', 'short-term',
'shrill', 'shy', 'sick', 'significant', 'silent', 'silky', 'silly',
'similar', 'simple', 'single', 'skilled', 'skinny', 'sleepy', 'slight',
'slim', 'slimy', 'slippery', 'slow', 'small', 'smart', 'smiling',
'smoggy', 'smooth', 'so-called', 'social', 'socialist', 'soft', 'solar',
'sole', 'solid', 'sophisticated', 'sore', 'sorry', 'sound', 'sour',
'southern', 'soviet', 'spanish', 'spare', 'sparkling', 'spatial',
'special', 'specific', 'specified', 'spectacular', 'spicy', 'spiritual',
'splendid', 'spontaneous', 'sporting', 'spotless', 'spotty', 'square',
'squealing', 'stable', 'stale', 'standard', 'static', 'statistical',
'statutory', 'steady', 'steep', 'sticky', 'stiff', 'still', 'stingy',
'stormy', 'straight', 'straightforward', 'strange', 'strategic',
'strict', 'striking', 'striped', 'strong', 'structural', 'stuck',
'stupid', 'subjective', 'subsequent', 'substantial', 'subtle',
'successful', 'successive', 'sudden', 'sufficient', 'suitable',
'sunny', 'super', 'superb', 'superior', 'supporting', 'supposed',
'supreme', 'sure', 'surprised', 'surprising', 'surrounding',
'surviving', 'suspicious', 'sweet', 'swift', 'swiss', 'symbolic',
'sympathetic', 'systematic', 'tall', 'tame', 'tan', 'tart',
'tasteless', 'tasty', 'technical', 'technological', 'teenage',
'temporary', 'tender', 'tense', 'terrible', 'territorial', 'testy',
'then', 'theoretical', 'thick', 'thin', 'thirsty', 'thorough',
'thoughtful', 'thoughtless', 'thundering', 'tight', 'tiny', 'tired',
'top', 'tory', 'total', 'tough', 'toxic', 'traditional', 'tragic',
'tremendous', 'tricky', 'tropical', 'troubled', 'turkish', 'typical',
'ugliest', 'ugly', 'ultimate', 'unable', 'unacceptable', 'unaware',
'uncertain', 'unchanged', 'uncomfortable', 'unconscious', 'underground',
'underlying', 'unemployed', 'uneven', 'unexpected', 'unfair',
'unfortunate', 'unhappy', 'uniform', 'uninterested', 'unique', 'united',
'universal', 'unknown', 'unlikely', 'unnecessary', 'unpleasant',
'unsightly', 'unusual', 'unwilling', 'upper', 'upset', 'uptight',
'urban', 'urgent', 'used', 'useful', 'useless', 'usual', 'vague',
'valid', 'valuable', 'variable', 'varied', 'various', 'varying', 'vast',
'verbal', 'vertical', 'very', 'victorian', 'victorious', 'video-taped',
'violent', 'visible', 'visiting', 'visual', 'vital', 'vivacious',
'vivid', 'vocational', 'voiceless', 'voluntary', 'vulnerable',
'wandering', 'warm', 'wasteful', 'watery', 'weak', 'wealthy', 'weary',
'wee', 'weekly', 'weird', 'welcome', 'well', 'well-known', 'welsh',
'western', 'wet', 'whispering', 'white', 'whole', 'wicked', 'wide',
'wide-eyed', 'widespread', 'wild', 'willing', 'wise', 'witty',
'wonderful', 'wooden', 'working', 'working-class', 'worldwide',
'worried', 'worrying', 'worthwhile', 'worthy', 'written', 'wrong',
'yellow', 'young', 'yummy', 'zany', 'zealous']
b = ['abiding', 'accelerating', 'accepting', 'accomplishing', 'achieving',
'acquiring', 'acteding', 'activating', 'adapting', 'adding', 'addressing',
'administering', 'admiring', 'admiting', 'adopting', 'advising', 'affording',
'agreeing', 'alerting', 'alighting', 'allowing', 'altereding', 'amusing',
'analyzing', 'announcing', 'annoying', 'answering', 'anticipating',
'apologizing', 'appearing', 'applauding', 'applieding', 'appointing',
'appraising', 'appreciating', 'approving', 'arbitrating', 'arguing',
'arising', 'arranging', 'arresting', 'arriving', 'ascertaining', 'asking',
'assembling', 'assessing', 'assisting', 'assuring', 'attaching', 'attacking',
'attaining', 'attempting', 'attending', 'attracting', 'auditeding', 'avoiding',
'awaking', 'backing', 'baking', 'balancing', 'baning', 'banging', 'baring',
'bating', 'bathing', 'battling', 'bing', 'beaming', 'bearing', 'beating',
'becoming', 'beging', 'begining', 'behaving', 'beholding', 'belonging',
'bending', 'beseting', 'beting', 'biding', 'binding', 'biting', 'bleaching',
'bleeding', 'blessing', 'blinding', 'blinking', 'bloting', 'blowing',
'blushing', 'boasting', 'boiling', 'bolting', 'bombing', 'booking',
'boring', 'borrowing', 'bouncing', 'bowing', 'boxing', 'braking',
'branching', 'breaking', 'breathing', 'breeding', 'briefing', 'bringing',
'broadcasting', 'bruising', 'brushing', 'bubbling', 'budgeting', 'building',
'bumping', 'burning', 'bursting', 'burying', 'busting', 'buying', 'buzing',
'calculating', 'calling', 'camping', 'caring', 'carrying', 'carving',
'casting', 'cataloging', 'catching', 'causing', 'challenging', 'changing',
'charging', 'charting', 'chasing', 'cheating', 'checking', 'cheering',
'chewing', 'choking', 'choosing', 'choping', 'claiming', 'claping',
'clarifying', 'classifying', 'cleaning', 'clearing', 'clinging', 'cliping',
'closing', 'clothing', 'coaching', 'coiling', 'collecting', 'coloring',
'combing', 'coming', 'commanding', 'communicating', 'comparing', 'competing',
'compiling', 'complaining', 'completing', 'composing', 'computing',
'conceiving', 'concentrating', 'conceptualizing', 'concerning', 'concluding',
'conducting', 'confessing', 'confronting', 'confusing', 'connecting',
'conserving', 'considering', 'consisting', 'consolidating', 'constructing',
'consulting', 'containing', 'continuing', 'contracting', 'controling',
'converting', 'coordinating', 'copying', 'correcting', 'correlating',
'costing', 'coughing', 'counseling', 'counting', 'covering', 'cracking',
'crashing', 'crawling', 'creating', 'creeping', 'critiquing', 'crossing',
'crushing', 'crying', 'curing', 'curling', 'curving', 'cuting', 'cycling',
'daming', 'damaging', 'dancing', 'daring', 'dealing', 'decaying', 'deceiving',
'deciding', 'decorating', 'defining', 'delaying', 'delegating', 'delighting',
'delivering', 'demonstrating', 'depending', 'describing', 'deserting',
'deserving', 'designing', 'destroying', 'detailing', 'detecting',
'determining', 'developing', 'devising', 'diagnosing', 'diging',
'directing', 'disagreing', 'disappearing', 'disapproving', 'disarming',
'discovering', 'disliking', 'dispensing', 'displaying', 'disproving',
'dissecting', 'distributing', 'diving', 'diverting', 'dividing', 'doing',
'doubling', 'doubting', 'drafting', 'draging', 'draining', 'dramatizing',
'drawing', 'dreaming', 'dressing', 'drinking', 'driping', 'driving',
'dropping', 'drowning', 'druming', 'drying', 'dusting', 'dwelling',
'earning', 'eating', 'editeding', 'educating', 'eliminating',
'embarrassing', 'employing', 'emptying', 'enacteding', 'encouraging',
'ending', 'enduring', 'enforcing', 'engineering', 'enhancing',
'enjoying', 'enlisting', 'ensuring', 'entering', 'entertaining',
'escaping', 'establishing', 'estimating', 'evaluating', 'examining',
'exceeding', 'exciting', 'excusing', 'executing', 'exercising', 'exhibiting',
'existing', 'expanding', 'expecting', 'expediting', 'experimenting',
'explaining', 'exploding', 'expressing', 'extending', 'extracting',
'facing', 'facilitating', 'fading', 'failing', 'fancying', 'fastening',
'faxing', 'fearing', 'feeding', 'feeling', 'fencing', 'fetching', 'fighting',
'filing', 'filling', 'filming', 'finalizing', 'financing', 'finding',
'firing', 'fiting', 'fixing', 'flaping', 'flashing', 'fleing', 'flinging',
'floating', 'flooding', 'flowing', 'flowering', 'flying', 'folding',
'following', 'fooling', 'forbiding', 'forcing', 'forecasting', 'foregoing',
'foreseing', 'foretelling', 'forgeting', 'forgiving', 'forming',
'formulating', 'forsaking', 'framing', 'freezing', 'frightening', 'frying',
'gathering', 'gazing', 'generating', 'geting', 'giving', 'glowing', 'gluing',
'going', 'governing', 'grabing', 'graduating', 'grating', 'greasing', 'greeting',
'grinning', 'grinding', 'griping', 'groaning', 'growing', 'guaranteeing',
'guarding', 'guessing', 'guiding', 'hammering', 'handing', 'handling',
'handwriting', 'hanging', 'happening', 'harassing', 'harming', 'hating',
'haunting', 'heading', 'healing', 'heaping', 'hearing', 'heating', 'helping',
'hiding', 'hitting', 'holding', 'hooking', 'hoping', 'hopping', 'hovering',
'hugging', 'hmuming', 'hunting', 'hurrying', 'hurting', 'hypothesizing',
'identifying', 'ignoring', 'illustrating', 'imagining', 'implementing',
'impressing', 'improving', 'improvising', 'including', 'increasing',
'inducing', 'influencing', 'informing', 'initiating', 'injecting',
'injuring', 'inlaying', 'innovating', 'inputing', 'inspecting',
'inspiring', 'installing', 'instituting', 'instructing', 'insuring',
'integrating', 'intending', 'intensifying', 'interesting',
'interfering', 'interlaying', 'interpreting', 'interrupting',
'interviewing', 'introducing', 'inventing', 'inventorying',
'investigating', 'inviting', 'irritating', 'itching', 'jailing',
'jamming', 'jogging', 'joining', 'joking', 'judging', 'juggling', 'jumping',
'justifying', 'keeping', 'kepting', 'kicking', 'killing', 'kissing', 'kneeling',
'kniting', 'knocking', 'knotting', 'knowing', 'labeling', 'landing', 'lasting',
'laughing', 'launching', 'laying', 'leading', 'leaning', 'leaping', 'learning',
'leaving', 'lecturing', 'leding', 'lending', 'leting', 'leveling',
'licensing', 'licking', 'lying', 'lifteding', 'lighting', 'lightening',
'liking', 'listing', 'listening', 'living', 'loading', 'locating',
'locking', 'loging', 'longing', 'looking', 'losing', 'loving',
'maintaining', 'making', 'maning', 'managing', 'manipulating',
'manufacturing', 'mapping', 'marching', 'marking', 'marketing',
'marrying', 'matching', 'mating', 'mattering', 'meaning', 'measuring',
'meddling', 'mediating', 'meeting', 'melting', 'melting', 'memorizing',
'mending', 'mentoring', 'milking', 'mining', 'misleading', 'missing',
'misspelling', 'mistaking', 'misunderstanding', 'mixing', 'moaning',
'modeling', 'modifying', 'monitoring', 'mooring', 'motivating',
'mourning', 'moving', 'mowing', 'muddling', 'muging', 'multiplying',
'murdering', 'nailing', 'naming', 'navigating', 'needing', 'negotiating',
'nesting', 'noding', 'nominating', 'normalizing', 'noting', 'noticing',
'numbering', 'obeying', 'objecting', 'observing', 'obtaining', 'occuring',
'offending', 'offering', 'officiating', 'opening', 'operating', 'ordering',
'organizing', 'orienteding', 'originating', 'overcoming', 'overdoing',
'overdrawing', 'overflowing', 'overhearing', 'overtaking', 'overthrowing',
'owing', 'owning', 'packing', 'paddling', 'painting', 'parking', 'parting',
'participating', 'passing', 'pasting', 'pating', 'pausing', 'paying',
'pecking', 'pedaling', 'peeling', 'peeping', 'perceiving', 'perfecting',
'performing', 'permiting', 'persuading', 'phoning', 'photographing',
'picking', 'piloting', 'pinching', 'pining', 'pinpointing', 'pioneering',
'placing', 'planing', 'planting', 'playing', 'pleading', 'pleasing',
'plugging', 'pointing', 'poking', 'polishing', 'poping', 'possessing',
'posting', 'pouring', 'practicing', 'praiseding', 'praying', 'preaching',
'preceding', 'predicting', 'prefering', 'preparing', 'prescribing',
'presenting', 'preserving', 'preseting', 'presiding', 'pressing',
'pretending', 'preventing', 'pricking', 'printing', 'processing',
'procuring', 'producing', 'professing', 'programing', 'progressing',
'projecting', 'promising', 'promoting', 'proofreading', 'proposing',
'protecting', 'proving', 'providing', 'publicizing', 'pulling', 'pumping',
'punching', 'puncturing', 'punishing', 'purchasing', 'pushing', 'puting',
'qualifying', 'questioning', 'queuing', 'quiting', 'racing', 'radiating',
'raining', 'raising', 'ranking', 'rating', 'reaching', 'reading',
'realigning', 'realizing', 'reasoning', 'receiving', 'recognizing',
'recommending', 'reconciling', 'recording', 'recruiting', 'reducing',
'referring', 'reflecting', 'refusing', 'regreting', 'regulating',
'rehabilitating', 'reigning', 'reinforcing', 'rejecting', 'rejoicing',
'relating', 'relaxing', 'releasing', 'relying', 'remaining', 'remembering',
'reminding', 'removing', 'rendering', 'reorganizing', 'repairing',
'repeating', 'replacing', 'replying', 'reporting', 'representing',
'reproducing', 'requesting', 'rescuing', 'researching', 'resolving',
'responding', 'restoreding', 'restructuring', 'retiring', 'retrieving',
'returning', 'reviewing', 'revising', 'rhyming', 'riding', 'riding',
'ringing', 'rinsing', 'rising', 'risking', 'robing', 'rocking', 'rolling',
'roting', 'rubing', 'ruining', 'ruling', 'runing', 'rushing', 'sacking',
'sailing', 'satisfying', 'saving', 'sawing', 'saying', 'scaring',
'scattering', 'scheduling', 'scolding', 'scorching', 'scraping',
'scratching', 'screaming', 'screwing', 'scribbling', 'scrubing',
'sealing', 'searching', 'securing', 'seing', 'seeking', 'selecting',
'selling', 'sending', 'sensing', 'separating', 'serving', 'servicing',
'seting', 'settling', 'sewing', 'shading', 'shaking', 'shaping',
'sharing', 'shaving', 'shearing', 'sheding', 'sheltering', 'shining',
'shivering', 'shocking', 'shoing', 'shooting', 'shoping', 'showing',
'shrinking', 'shruging', 'shuting', 'sighing', 'signing', 'signaling',
'simplifying', 'sining', 'singing', 'sinking', 'siping', 'siting',
'sketching', 'skiing', 'skiping', 'slaping', 'slaying', 'sleeping',
'sliding', 'slinging', 'slinking', 'sliping', 'sliting', 'slowing',
'smashing', 'smelling', 'smiling', 'smiting', 'smoking', 'snatching',
'sneaking', 'sneezing', 'sniffing', 'snoring', 'snowing', 'soaking',
'solving', 'soothing', 'soothsaying', 'sorting', 'sounding', 'sowing',
'sparing', 'sparking', 'sparkling', 'speaking', 'specifying', 'speeding',
'spelling', 'spending', 'spilling', 'spining', 'spiting', 'spliting',
'spoiling', 'spoting', 'spraying', 'spreading', 'springing', 'sprouting',
'squashing', 'squeaking', 'squealing', 'squeezing', 'staining', 'stamping',
'standing', 'staring', 'starting', 'staying', 'stealing', 'steering',
'stepping', 'sticking', 'stimulating', 'stinging', 'stinking', 'stirring',
'stitching', 'stoping', 'storing', 'straping', 'streamlining',
'strengthening', 'stretching', 'striding', 'striking', 'stringing',
'stripping', 'striving', 'stroking', 'structuring', 'studying',
'stuffing', 'subleting', 'subtracting', 'succeeding', 'sucking',
'suffering', 'suggesting', 'suiting', 'summarizing', 'supervising',
'supplying', 'supporting', 'supposing', 'surprising', 'surrounding',
'suspecting', 'suspending', 'swearing', 'sweating', 'sweeping', 'swelling',
'swimming', 'swinging', 'switching', 'symbolizing', 'synthesizing',
'systemizing', 'tabulating', 'taking', 'talking', 'taming', 'taping',
'targeting', 'tasting', 'teaching', 'tearing', 'teasing', 'telephoning',
'telling', 'tempting', 'terrifying', 'testing', 'thanking', 'thawing',
'thinking', 'thriving', 'throwing', 'thrusting', 'ticking', 'tickling',
'tying', 'timing', 'tiping', 'tiring', 'touching', 'touring', 'towing',
'tracing', 'trading', 'training', 'transcribing', 'transfering',
'transforming', 'translating', 'transporting', 'traping', 'traveling',
'treading', 'treating', 'trembling', 'tricking', 'triping', 'troting',
'troubling', 'troubleshooting', 'trusting', 'trying', 'tuging', 'tumbling',
'turning', 'tutoring', 'twisting', 'typing', 'undergoing', 'understanding',
'undertaking', 'undressing', 'unfastening', 'unifying', 'uniting',
'unlocking', 'unpacking', 'untidying', 'updating', 'upgrading',
'upholding', 'upseting', 'using', 'utilizing', 'vanishing', 'verbalizing',
'verifying', 'vexing', 'visiting', 'wailing', 'waiting', 'waking',
'walking', 'wandering', 'wanting', 'warming', 'warning', 'washing',
'wasting', 'watching', 'watering', 'waving', 'wearing', 'weaving',
'wedding', 'weeping', 'weighing', 'welcoming', 'wending', 'weting',
'whining', 'whiping', 'whirling', 'whispering', 'whistling', 'wining',
'winding', 'winking', 'wiping', 'wishing', 'withdrawing', 'withholding',
'withstanding', 'wobbling', 'wondering', 'working', 'worrying', 'wrapping',
'wrecking', 'wrestling', 'wriggling', 'wringing', 'writing', 'x-raying',
'yawning', 'yelling', 'zipping', 'zooming']
================================================
FILE: BigGAN_utils/binary_utils.py
================================================
from torch.autograd import Function
from torch.optim import SGD
class BinaryActivation(Function):
@staticmethod
def forward(ctx, x):
ctx.save_for_backward(x)
return (x.sign() + 1.) / 2.
@staticmethod
def backward(ctx, grad_output):
return grad_output.clone()
================================================
FILE: BigGAN_utils/calculate_inception_moments.py
================================================
''' Calculate Inception Moments
This script iterates over the dataset and calculates the moments of the
activations of the Inception net (needed for FID), and also returns
the Inception Score of the training data.
Note that if you don't shuffle the data, the IS of true data will be under-
estimated as it is label-ordered. By default, the data is not shuffled
so as to reduce non-determinism. '''
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import utils
import inception_utils
from tqdm import tqdm, trange
from argparse import ArgumentParser
def prepare_parser():
usage = 'Calculate and store inception metrics.'
parser = ArgumentParser(description=usage)
parser.add_argument(
'--dataset', type=str, default='I128_hdf5',
help='Which Dataset to train on, out of I128, I256, C10, C100...'
'Append _hdf5 to use the hdf5 version of the dataset. (default: %(default)s)')
parser.add_argument(
'--data_root', type=str, default='data',
help='Default location where data is stored (default: %(default)s)')
parser.add_argument(
'--batch_size', type=int, default=64,
help='Default overall batchsize (default: %(default)s)')
parser.add_argument(
'--parallel', action='store_true', default=False,
help='Train with multiple GPUs (default: %(default)s)')
parser.add_argument(
'--augment', action='store_true', default=False,
help='Augment with random crops and flips (default: %(default)s)')
parser.add_argument(
'--num_workers', type=int, default=8,
help='Number of dataloader workers (default: %(default)s)')
parser.add_argument(
'--shuffle', action='store_true', default=False,
help='Shuffle the data? (default: %(default)s)')
parser.add_argument(
'--seed', type=int, default=0,
help='Random seed to use.')
return parser
def run(config):
# Get loader
config['drop_last'] = False
loaders = utils.get_data_loaders(**config)
# Load inception net
net = inception_utils.load_inception_net(parallel=config['parallel'])
pool, logits, labels = [], [], []
device = 'cuda'
for i, (x, y) in enumerate(tqdm(loaders[0])):
x = x.to(device)
with torch.no_grad():
pool_val, logits_val = net(x)
pool += [np.asarray(pool_val.cpu())]
logits += [np.asarray(F.softmax(logits_val, 1).cpu())]
labels += [np.asarray(y.cpu())]
pool, logits, labels = [np.concatenate(item, 0) for item in [pool, logits, labels]]
# uncomment to save pool, logits, and labels to disk
# print('Saving pool, logits, and labels to disk...')
# np.savez(config['dataset']+'_inception_activations.npz',
# {'pool': pool, 'logits': logits, 'labels': labels})
# Calculate inception metrics and report them
print('Calculating inception metrics...')
IS_mean, IS_std = inception_utils.calculate_inception_score(logits)
print('Training data from dataset %s has IS of %5.5f +/- %5.5f' % (config['dataset'], IS_mean, IS_std))
# Prepare mu and sigma, save to disk. Remove "hdf5" by default
# (the FID code also knows to strip "hdf5")
print('Calculating means and covariances...')
mu, sigma = np.mean(pool, axis=0), np.cov(pool, rowvar=False)
print('Saving calculated means and covariances to disk...')
np.savez(config['dataset'].strip('_hdf5')+'_inception_moments.npz', **{'mu' : mu, 'sigma' : sigma})
def main():
# parse command line
parser = prepare_parser()
config = vars(parser.parse_args())
print(config)
run(config)
if __name__ == '__main__':
main()
================================================
FILE: BigGAN_utils/datasets.py
================================================
''' Datasets
This file contains definitions for our CIFAR, ImageFolder, and HDF5 datasets
'''
import os
import os.path
import sys
from PIL import Image
import numpy as np
from tqdm import tqdm, trange
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torchvision.datasets.utils import download_url, check_integrity
import torch.utils.data as data
from torch.utils.data import DataLoader
IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm']
def is_image_file(filename):
"""Checks if a file is an image.
Args:
filename (string): path to a file
Returns:
bool: True if the filename ends with a known image extension
"""
filename_lower = filename.lower()
return any(filename_lower.endswith(ext) for ext in IMG_EXTENSIONS)
def find_classes(dir):
classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))]
classes.sort()
class_to_idx = {classes[i]: i for i in range(len(classes))}
return classes, class_to_idx
def make_dataset(dir, class_to_idx):
images = []
dir = os.path.expanduser(dir)
for target in tqdm(sorted(os.listdir(dir))):
d = os.path.join(dir, target)
if not os.path.isdir(d):
continue
for root, _, fnames in sorted(os.walk(d)):
for fname in sorted(fnames):
if is_image_file(fname):
path = os.path.join(root, fname)
item = (path, class_to_idx[target])
images.append(item)
return images
def pil_loader(path):
# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
with open(path, 'rb') as f:
img = Image.open(f)
return img.convert('RGB')
def accimage_loader(path):
import accimage
try:
return accimage.Image(path)
except IOError:
# Potentially a decoding problem, fall back to PIL.Image
return pil_loader(path)
def default_loader(path):
from torchvision import get_image_backend
if get_image_backend() == 'accimage':
return accimage_loader(path)
else:
return pil_loader(path)
class ImageFolder(data.Dataset):
"""A generic data loader where the images are arranged in this way: ::
root/dogball/xxx.png
root/dogball/xxy.png
root/dogball/xxz.png
root/cat/123.png
root/cat/nsdf3.png
root/cat/asd932_.png
Args:
root (string): Root directory path.
transform (callable, optional): A function/transform that takes in an PIL image
and returns a transformed version. E.g, ``transforms.RandomCrop``
target_transform (callable, optional): A function/transform that takes in the
target and transforms it.
loader (callable, optional): A function to load an image given its path.
Attributes:
classes (list): List of the class names.
class_to_idx (dict): Dict with items (class_name, class_index).
imgs (list): List of (image path, class_index) tuples
"""
def __init__(self, root, transform=None, target_transform=None,
loader=default_loader, load_in_mem=False,
index_filename='imagenet_imgs.npz', **kwargs):
classes, class_to_idx = find_classes(root)
# Load pre-computed image directory walk
if os.path.exists(index_filename):
print('Loading pre-saved Index file %s...' % index_filename)
imgs = np.load(index_filename)['imgs']
# If first time, walk the folder directory and save the
# results to a pre-computed file.
else:
print('Generating Index file %s...' % index_filename)
imgs = make_dataset(root, class_to_idx)
np.savez_compressed(index_filename, **{'imgs' : imgs})
if len(imgs) == 0:
raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n"
"Supported image extensions are: " + ",".join(IMG_EXTENSIONS)))
self.root = root
self.imgs = imgs
self.classes = classes
self.class_to_idx = class_to_idx
self.transform = transform
self.target_transform = target_transform
self.loader = loader
self.load_in_mem = load_in_mem
if self.load_in_mem:
print('Loading all images into memory...')
self.data, self.labels = [], []
for index in tqdm(range(len(self.imgs))):
path, target = imgs[index][0], imgs[index][1]
self.data.append(self.transform(self.loader(path)))
self.labels.append(target)
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: (image, target) where target is class_index of the target class.
"""
if self.load_in_mem:
img = self.data[index]
target = self.labels[index]
else:
path, target = self.imgs[index]
img = self.loader(str(path))
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
# print(img.size(), target)
return img, int(target)
def __len__(self):
return len(self.imgs)
def __repr__(self):
fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
fmt_str += ' Number of datapoints: {}\n'.format(self.__len__())
fmt_str += ' Root Location: {}\n'.format(self.root)
tmp = ' Transforms (if any): '
fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
tmp = ' Target Transforms (if any): '
fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
return fmt_str
''' ILSVRC_HDF5: A dataset to support I/O from an HDF5 to avoid
having to load individual images all the time. '''
import h5py as h5
import torch
class ILSVRC_HDF5(data.Dataset):
def __init__(self, root, transform=None, target_transform=None,
load_in_mem=False, train=True,download=False, validate_seed=0,
val_split=0, **kwargs): # last four are dummies
self.root = root
self.num_imgs = len(h5.File(root, 'r')['labels'])
# self.transform = transform
self.target_transform = target_transform
# Set the transform here
self.transform = transform
# load the entire dataset into memory?
self.load_in_mem = load_in_mem
# If loading into memory, do so now
if self.load_in_mem:
print('Loading %s into memory...' % root)
with h5.File(root,'r') as f:
self.data = f['imgs'][:]
self.labels = f['labels'][:]
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: (image, target) where target is class_index of the target class.
"""
# If loaded the entire dataset in RAM, get image from memory
if self.load_in_mem:
img = self.data[index]
target = self.labels[index]
# Else load it from disk
else:
with h5.File(self.root,'r') as f:
img = f['imgs'][index]
target = f['labels'][index]
# if self.transform is not None:
# img = self.transform(img)
# Apply my own transform
img = ((torch.from_numpy(img).float() / 255) - 0.5) * 2
if self.target_transform is not None:
target = self.target_transform(target)
return img, int(target)
def __len__(self):
return self.num_imgs
# return len(self.f['imgs'])
import pickle
class CIFAR10(dset.CIFAR10):
def __init__(self, root, train=True,
transform=None, target_transform=None,
download=True, validate_seed=0,
val_split=0, load_in_mem=True, **kwargs):
self.root = os.path.expanduser(root)
self.transform = transform
self.target_transform = target_transform
self.train = train # training set or test set
self.val_split = val_split
if download:
self.download()
if not self._check_integrity():
raise RuntimeError('Dataset not found or corrupted.' +
' You can use download=True to download it')
# now load the picked numpy arrays
self.data = []
self.labels= []
for fentry in self.train_list:
f = fentry[0]
file = os.path.join(self.root, self.base_folder, f)
fo = open(file, 'rb')
if sys.version_info[0] == 2:
entry = pickle.load(fo)
else:
entry = pickle.load(fo, encoding='latin1')
self.data.append(entry['data'])
if 'labels' in entry:
self.labels += entry['labels']
else:
self.labels += entry['fine_labels']
fo.close()
self.data = np.concatenate(self.data)
# Randomly select indices for validation
if self.val_split > 0:
label_indices = [[] for _ in range(max(self.labels)+1)]
for i,l in enumerate(self.labels):
label_indices[l] += [i]
label_indices = np.asarray(label_indices)
# randomly grab 500 elements of each class
np.random.seed(validate_seed)
self.val_indices = []
for l_i in label_indices:
self.val_indices += list(l_i[np.random.choice(len(l_i), int(len(self.data) * val_split) // (max(self.labels) + 1) ,replace=False)])
if self.train=='validate':
self.data = self.data[self.val_indices]
self.labels = list(np.asarray(self.labels)[self.val_indices])
self.data = self.data.reshape((int(50e3 * self.val_split), 3, 32, 32))
self.data = self.data.transpose((0, 2, 3, 1)) # convert to HWC
elif self.train:
print(np.shape(self.data))
if self.val_split > 0:
self.data = np.delete(self.data,self.val_indices,axis=0)
self.labels = list(np.delete(np.asarray(self.labels),self.val_indices,axis=0))
self.data = self.data.reshape((int(50e3 * (1.-self.val_split)), 3, 32, 32))
self.data = self.data.transpose((0, 2, 3, 1)) # convert to HWC
else:
f = self.test_list[0][0]
file = os.path.join(self.root, self.base_folder, f)
fo = open(file, 'rb')
if sys.version_info[0] == 2:
entry = pickle.load(fo)
else:
entry = pickle.load(fo, encoding='latin1')
self.data = entry['data']
if 'labels' in entry:
self.labels = entry['labels']
else:
self.labels = entry['fine_labels']
fo.close()
self.data = self.data.reshape((10000, 3, 32, 32))
self.data = self.data.transpose((0, 2, 3, 1)) # convert to HWC
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: (image, target) where target is index of the target class.
"""
img, target = self.data[index], self.labels[index]
# doing this so that it is consistent with all other datasets
# to return a PIL Image
img = Image.fromarray(img)
if self.transform is not None:
img = self.transform(img)
if self.target_transform is not None:
target = self.target_transform(target)
return img, target
def __len__(self):
return len(self.data)
class CIFAR100(CIFAR10):
base_folder = 'cifar-100-python'
url = "http://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz"
filename = "cifar-100-python.tar.gz"
tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85'
train_list = [
['train', '16019d7e3df5f24257cddd939b257f8d'],
]
test_list = [
['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'],
]
================================================
FILE: BigGAN_utils/inception_tf13.py
================================================
''' Tensorflow inception score code
Derived from https://github.com/openai/improved-gan
Code derived from tensorflow/tensorflow/models/image/imagenet/classify_image.py
THIS CODE REQUIRES TENSORFLOW 1.3 or EARLIER to run in PARALLEL BATCH MODE
To use this code, run sample.py on your model with --sample_npz, and then
pass the experiment name in the --experiment_name.
This code also saves pool3 stats to an npz file for FID calculation
'''
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import os.path
import sys
import tarfile
import math
from tqdm import tqdm, trange
from argparse import ArgumentParser
import numpy as np
from six.moves import urllib
import tensorflow as tf
MODEL_DIR = ''
DATA_URL = 'http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz'
softmax = None
def prepare_parser():
usage = 'Parser for TF1.3- Inception Score scripts.'
parser = ArgumentParser(description=usage)
parser.add_argument(
'--experiment_name', type=str, default='',
help='Which experiment''s samples.npz file to pull and evaluate')
parser.add_argument(
'--experiment_root', type=str, default='samples',
help='Default location where samples are stored (default: %(default)s)')
parser.add_argument(
'--batch_size', type=int, default=500,
help='Default overall batchsize (default: %(default)s)')
return parser
def run(config):
# Inception with TF1.3 or earlier.
# Call this function with list of images. Each of elements should be a
# numpy array with values ranging from 0 to 255.
def get_inception_score(images, splits=10):
assert(type(images) == list)
assert(type(images[0]) == np.ndarray)
assert(len(images[0].shape) == 3)
assert(np.max(images[0]) > 10)
assert(np.min(images[0]) >= 0.0)
inps = []
for img in images:
img = img.astype(np.float32)
inps.append(np.expand_dims(img, 0))
bs = config['batch_size']
with tf.Session() as sess:
preds, pools = [], []
n_batches = int(math.ceil(float(len(inps)) / float(bs)))
for i in trange(n_batches):
inp = inps[(i * bs):min((i + 1) * bs, len(inps))]
inp = np.concatenate(inp, 0)
pred, pool = sess.run([softmax, pool3], {'ExpandDims:0': inp})
preds.append(pred)
pools.append(pool)
preds = np.concatenate(preds, 0)
scores = []
for i in range(splits):
part = preds[(i * preds.shape[0] // splits):((i + 1) * preds.shape[0] // splits), :]
kl = part * (np.log(part) - np.log(np.expand_dims(np.mean(part, 0), 0)))
kl = np.mean(np.sum(kl, 1))
scores.append(np.exp(kl))
return np.mean(scores), np.std(scores), np.squeeze(np.concatenate(pools, 0))
# Init inception
def _init_inception():
global softmax, pool3
if not os.path.exists(MODEL_DIR):
os.makedirs(MODEL_DIR)
filename = DATA_URL.split('/')[-1]
filepath = os.path.join(MODEL_DIR, filename)
if not os.path.exists(filepath):
def _progress(count, block_size, total_size):
sys.stdout.write('\r>> Downloading %s %.1f%%' % (
filename, float(count * block_size) / float(total_size) * 100.0))
sys.stdout.flush()
filepath, _ = urllib.request.urlretrieve(DATA_URL, filepath, _progress)
print()
statinfo = os.stat(filepath)
print('Succesfully downloaded', filename, statinfo.st_size, 'bytes.')
tarfile.open(filepath, 'r:gz').extractall(MODEL_DIR)
with tf.gfile.FastGFile(os.path.join(
MODEL_DIR, 'classify_image_graph_def.pb'), 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
_ = tf.import_graph_def(graph_def, name='')
# Works with an arbitrary minibatch size.
with tf.Session() as sess:
pool3 = sess.graph.get_tensor_by_name('pool_3:0')
ops = pool3.graph.get_operations()
for op_idx, op in enumerate(ops):
for o in op.outputs:
shape = o.get_shape()
shape = [s.value for s in shape]
new_shape = []
for j, s in enumerate(shape):
if s == 1 and j == 0:
new_shape.append(None)
else:
new_shape.append(s)
o._shape = tf.TensorShape(new_shape)
w = sess.graph.get_operation_by_name("softmax/logits/MatMul").inputs[1]
logits = tf.matmul(tf.squeeze(pool3), w)
softmax = tf.nn.softmax(logits)
# if softmax is None: # No need to functionalize like this.
_init_inception()
fname = '%s/%s/samples.npz' % (config['experiment_root'], config['experiment_name'])
print('loading %s ...'%fname)
ims = np.load(fname)['x']
import time
t0 = time.time()
inc_mean, inc_std, pool_activations = get_inception_score(list(ims.swapaxes(1,2).swapaxes(2,3)), splits=10)
t1 = time.time()
print('Saving pool to numpy file for FID calculations...')
np.savez('%s/%s/TF_pool.npz' % (config['experiment_root'], config['experiment_name']), **{'pool_mean': np.mean(pool_activations,axis=0), 'pool_var': np.cov(pool_activations, rowvar=False)})
print('Inception took %3f seconds, score of %3f +/- %3f.'%(t1-t0, inc_mean, inc_std))
def main():
# parse command line and run
parser = prepare_parser()
config = vars(parser.parse_args())
print(config)
run(config)
if __name__ == '__main__':
main()
================================================
FILE: BigGAN_utils/inception_utils.py
================================================
''' Inception utilities
This file contains methods for calculating IS and FID, using either
the original numpy code or an accelerated fully-pytorch version that
uses a fast newton-schulz approximation for the matrix sqrt. There are also
methods for acquiring a desired number of samples from the Generator,
and parallelizing the inbuilt PyTorch inception network.
NOTE that Inception Scores and FIDs calculated using these methods will
*not* be directly comparable to values calculated using the original TF
IS/FID code. You *must* use the TF model if you wish to report and compare
numbers. This code tends to produce IS values that are 5-10% lower than
those obtained through TF.
'''
import numpy as np
from scipy import linalg # For numpy FID
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import Parameter as P
from torchvision.models.inception import inception_v3
# Module that wraps the inception network to enable use with dataparallel and
# returning pool features and logits.
class WrapInception(nn.Module):
def __init__(self, net):
super(WrapInception,self).__init__()
self.net = net
self.mean = P(torch.tensor([0.485, 0.456, 0.406]).view(1, -1, 1, 1),
requires_grad=False)
self.std = P(torch.tensor([0.229, 0.224, 0.225]).view(1, -1, 1, 1),
requires_grad=False)
def forward(self, x):
# Normalize x
x = (x + 1.) / 2.0
x = (x - self.mean) / self.std
# Upsample if necessary
if x.shape[2] != 299 or x.shape[3] != 299:
x = F.interpolate(x, size=(299, 299), mode='bilinear', align_corners=True)
# 299 x 299 x 3
x = self.net.Conv2d_1a_3x3(x)
# 149 x 149 x 32
x = self.net.Conv2d_2a_3x3(x)
# 147 x 147 x 32
x = self.net.Conv2d_2b_3x3(x)
# 147 x 147 x 64
x = F.max_pool2d(x, kernel_size=3, stride=2)
# 73 x 73 x 64
x = self.net.Conv2d_3b_1x1(x)
# 73 x 73 x 80
x = self.net.Conv2d_4a_3x3(x)
# 71 x 71 x 192
x = F.max_pool2d(x, kernel_size=3, stride=2)
# 35 x 35 x 192
x = self.net.Mixed_5b(x)
# 35 x 35 x 256
x = self.net.Mixed_5c(x)
# 35 x 35 x 288
x = self.net.Mixed_5d(x)
# 35 x 35 x 288
x = self.net.Mixed_6a(x)
# 17 x 17 x 768
x = self.net.Mixed_6b(x)
# 17 x 17 x 768
x = self.net.Mixed_6c(x)
# 17 x 17 x 768
x = self.net.Mixed_6d(x)
# 17 x 17 x 768
x = self.net.Mixed_6e(x)
# 17 x 17 x 768
# 17 x 17 x 768
x = self.net.Mixed_7a(x)
# 8 x 8 x 1280
x = self.net.Mixed_7b(x)
# 8 x 8 x 2048
x = self.net.Mixed_7c(x)
# 8 x 8 x 2048
pool = torch.mean(x.view(x.size(0), x.size(1), -1), 2)
# 1 x 1 x 2048
logits = self.net.fc(F.dropout(pool, training=False).view(pool.size(0), -1))
# 1000 (num_classes)
return pool, logits
# A pytorch implementation of cov, from Modar M. Alfadly
# https://discuss.pytorch.org/t/covariance-and-gradient-support/16217/2
def torch_cov(m, rowvar=False):
'''Estimate a covariance matrix given data.
Covariance indicates the level to which two variables vary together.
If we examine N-dimensional samples, `X = [x_1, x_2, ... x_N]^T`,
then the covariance matrix element `C_{ij}` is the covariance of
`x_i` and `x_j`. The element `C_{ii}` is the variance of `x_i`.
Args:
m: A 1-D or 2-D array containing multiple variables and observations.
Each row of `m` represents a variable, and each column a single
observation of all those variables.
rowvar: If `rowvar` is True, then each row represents a
variable, with observations in the columns. Otherwise, the
relationship is transposed: each column represents a variable,
while the rows contain observations.
Returns:
The covariance matrix of the variables.
'''
if m.dim() > 2:
raise ValueError('m has more than 2 dimensions')
if m.dim() < 2:
m = m.view(1, -1)
if not rowvar and m.size(0) != 1:
m = m.t()
# m = m.type(torch.double) # uncomment this line if desired
fact = 1.0 / (m.size(1) - 1)
m -= torch.mean(m, dim=1, keepdim=True)
mt = m.t() # if complex: mt = m.t().conj()
return fact * m.matmul(mt).squeeze()
# Pytorch implementation of matrix sqrt, from Tsung-Yu Lin, and Subhransu Maji
# https://github.com/msubhransu/matrix-sqrt
def sqrt_newton_schulz(A, numIters, dtype=None):
with torch.no_grad():
if dtype is None:
dtype = A.type()
batchSize = A.shape[0]
dim = A.shape[1]
normA = A.mul(A).sum(dim=1).sum(dim=1).sqrt()
Y = A.div(normA.view(batchSize, 1, 1).expand_as(A));
I = torch.eye(dim,dim).view(1, dim, dim).repeat(batchSize,1,1).type(dtype)
Z = torch.eye(dim,dim).view(1, dim, dim).repeat(batchSize,1,1).type(dtype)
for i in range(numIters):
T = 0.5*(3.0*I - Z.bmm(Y))
Y = Y.bmm(T)
Z = T.bmm(Z)
sA = Y*torch.sqrt(normA).view(batchSize, 1, 1).expand_as(A)
return sA
# FID calculator from TTUR--consider replacing this with GPU-accelerated cov
# calculations using torch?
def numpy_calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
"""Numpy implementation of the Frechet Distance.
Taken from https://github.com/bioinf-jku/TTUR
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
and X_2 ~ N(mu_2, C_2) is
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
Stable version by Dougal J. Sutherland.
Params:
-- mu1 : Numpy array containing the activations of a layer of the
inception net (like returned by the function 'get_predictions')
for generated samples.
-- mu2 : The sample mean over activations, precalculated on an
representive data set.
-- sigma1: The covariance matrix over activations for generated samples.
-- sigma2: The covariance matrix over activations, precalculated on an
representive data set.
Returns:
-- : The Frechet Distance.
"""
mu1 = np.atleast_1d(mu1)
mu2 = np.atleast_1d(mu2)
sigma1 = np.atleast_2d(sigma1)
sigma2 = np.atleast_2d(sigma2)
assert mu1.shape == mu2.shape, \
'Training and test mean vectors have different lengths'
assert sigma1.shape == sigma2.shape, \
'Training and test covariances have different dimensions'
diff = mu1 - mu2
# Product might be almost singular
covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
if not np.isfinite(covmean).all():
msg = ('fid calculation produces singular product; '
'adding %s to diagonal of cov estimates') % eps
print(msg)
offset = np.eye(sigma1.shape[0]) * eps
covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
# Numerical error might give slight imaginary component
if np.iscomplexobj(covmean):
print('wat')
if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
m = np.max(np.abs(covmean.imag))
raise ValueError('Imaginary component {}'.format(m))
covmean = covmean.real
tr_covmean = np.trace(covmean)
out = diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean
return out
def torch_calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
"""Pytorch implementation of the Frechet Distance.
Taken from https://github.com/bioinf-jku/TTUR
The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
and X_2 ~ N(mu_2, C_2) is
d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
Stable version by Dougal J. Sutherland.
Params:
-- mu1 : Numpy array containing the activations of a layer of the
inception net (like returned by the function 'get_predictions')
for generated samples.
-- mu2 : The sample mean over activations, precalculated on an
representive data set.
-- sigma1: The covariance matrix over activations for generated samples.
-- sigma2: The covariance matrix over activations, precalculated on an
representive data set.
Returns:
-- : The Frechet Distance.
"""
assert mu1.shape == mu2.shape, \
'Training and test mean vectors have different lengths'
assert sigma1.shape == sigma2.shape, \
'Training and test covariances have different dimensions'
diff = mu1 - mu2
# Run 50 itrs of newton-schulz to get the matrix sqrt of sigma1 dot sigma2
covmean = sqrt_newton_schulz(sigma1.mm(sigma2).unsqueeze(0), 50).squeeze()
out = (diff.dot(diff) + torch.trace(sigma1) + torch.trace(sigma2)
- 2 * torch.trace(covmean))
return out
# Calculate Inception Score mean + std given softmax'd logits and number of splits
def calculate_inception_score(pred, num_splits=10):
scores = []
for index in range(num_splits):
pred_chunk = pred[index * (pred.shape[0] // num_splits): (index + 1) * (pred.shape[0] // num_splits), :]
kl_inception = pred_chunk * (np.log(pred_chunk) - np.log(np.expand_dims(np.mean(pred_chunk, 0), 0)))
kl_inception = np.mean(np.sum(kl_inception, 1))
scores.append(np.exp(kl_inception))
return np.mean(scores), np.std(scores)
# Loop and run the sampler and the net until it accumulates num_inception_images
# activations. Return the pool, the logits, and the labels (if one wants
# Inception Accuracy the labels of the generated class will be needed)
def accumulate_inception_activations(sample, net, num_inception_images=50000):
pool, logits, labels = [], [], []
while (torch.cat(logits, 0).shape[0] if len(logits) else 0) < num_inception_images:
with torch.no_grad():
images, labels_val = sample()
pool_val, logits_val = net(images.float())
pool += [pool_val]
logits += [F.softmax(logits_val, 1)]
labels += [labels_val]
return torch.cat(pool, 0), torch.cat(logits, 0), torch.cat(labels, 0)
# Load and wrap the Inception model
def load_inception_net(parallel=False):
inception_model = inception_v3(pretrained=True, transform_input=False)
inception_model = WrapInception(inception_model.eval()).cuda()
if parallel:
print('Parallelizing Inception module...')
inception_model = nn.DataParallel(inception_model)
return inception_model
# This produces a function which takes in an iterator which returns a set number of samples
# and iterates until it accumulates config['num_inception_images'] images.
# The iterator can return samples with a different batch size than used in
# training, using the setting confg['inception_batchsize']
def prepare_inception_metrics(dataset, parallel, no_fid=False):
# Load metrics; this is intentionally not in a try-except loop so that
# the script will crash here if it cannot find the Inception moments.
# By default, remove the "hdf5" from dataset
dataset = dataset.strip('_hdf5')
data_mu = np.load(dataset+'_inception_moments.npz')['mu']
data_sigma = np.load(dataset+'_inception_moments.npz')['sigma']
# Load network
net = load_inception_net(parallel)
def get_inception_metrics(sample, num_inception_images, num_splits=10,
prints=True, use_torch=True):
if prints:
print('Gathering activations...')
pool, logits, labels = accumulate_inception_activations(sample, net, num_inception_images)
if prints:
print('Calculating Inception Score...')
IS_mean, IS_std = calculate_inception_score(logits.cpu().numpy(), num_splits)
if no_fid:
FID = 9999.0
else:
if prints:
print('Calculating means and covariances...')
if use_torch:
mu, sigma = torch.mean(pool, 0), torch_cov(pool, rowvar=False)
else:
mu, sigma = np.mean(pool.cpu().numpy(), axis=0), np.cov(pool.cpu().numpy(), rowvar=False)
if prints:
print('Covariances calculated, getting FID...')
if use_torch:
FID = torch_calculate_frechet_distance(mu, sigma, torch.tensor(data_mu).float().cuda(), torch.tensor(data_sigma).float().cuda())
FID = float(FID.cpu().numpy())
else:
FID = numpy_calculate_frechet_distance(mu.cpu().numpy(), sigma.cpu().numpy(), data_mu, data_sigma)
# Delete mu, sigma, pool, logits, and labels, just in case
del mu, sigma, pool, logits, labels
return IS_mean, IS_std, FID
return get_inception_metrics
================================================
FILE: BigGAN_utils/layers.py
================================================
''' Layers
This file contains various layers for the BigGAN models.
'''
import numpy as np
import torch
import torch.nn as nn
from torch.nn import init
import torch.optim as optim
import torch.nn.functional as F
from torch.nn import Parameter as P
from sync_batchnorm import SynchronizedBatchNorm2d as SyncBN2d
# Projection of x onto y
def proj(x, y):
return torch.mm(y, x.t()) * y / torch.mm(y, y.t())
# Orthogonalize x wrt list of vectors ys
def gram_schmidt(x, ys):
for y in ys:
x = x - proj(x, y)
return x
# Apply num_itrs steps of the power method to estimate top N singular values.
def power_iteration(W, u_, update=True, eps=1e-12):
# Lists holding singular vectors and values
us, vs, svs = [], [], []
for i, u in enumerate(u_):
# Run one step of the power iteration
with torch.no_grad():
v = torch.matmul(u, W)
# Run Gram-Schmidt to subtract components of all other singular vectors
v = F.normalize(gram_schmidt(v, vs), eps=eps)
# Add to the list
vs += [v]
# Update the other singular vector
u = torch.matmul(v, W.t())
# Run Gram-Schmidt to subtract components of all other singular vectors
u = F.normalize(gram_schmidt(u, us), eps=eps)
# Add to the list
us += [u]
if update:
u_[i][:] = u
# Compute this singular value and add it to the list
svs += [torch.squeeze(torch.matmul(torch.matmul(v, W.t()), u.t()))]
#svs += [torch.sum(F.linear(u, W.transpose(0, 1)) * v)]
return svs, us, vs
# Convenience passthrough function
class identity(nn.Module):
def forward(self, input):
return input
# Spectral normalization base class
class SN(object):
def __init__(self, num_svs, num_itrs, num_outputs, transpose=False, eps=1e-12):
# Number of power iterations per step
self.num_itrs = num_itrs
# Number of singular values
self.num_svs = num_svs
# Transposed?
self.transpose = transpose
# Epsilon value for avoiding divide-by-0
self.eps = eps
# Register a singular vector for each sv
for i in range(self.num_svs):
self.register_buffer('u%d' % i, torch.randn(1, num_outputs))
self.register_buffer('sv%d' % i, torch.ones(1))
# Singular vectors (u side)
@property
def u(self):
return [getattr(self, 'u%d' % i) for i in range(self.num_svs)]
# Singular values;
# note that these buffers are just for logging and are not used in training.
@property
def sv(self):
return [getattr(self, 'sv%d' % i) for i in range(self.num_svs)]
# Compute the spectrally-normalized weight
def W_(self):
W_mat = self.weight.view(self.weight.size(0), -1)
if self.transpose:
W_mat = W_mat.t()
# Apply num_itrs power iterations
for _ in range(self.num_itrs):
svs, us, vs = power_iteration(W_mat, self.u, update=self.training, eps=self.eps)
# Update the svs
if self.training:
with torch.no_grad(): # Make sure to do this in a no_grad() context or you'll get memory leaks!
for i, sv in enumerate(svs):
self.sv[i][:] = sv
return self.weight / svs[0]
# 2D Conv layer with spectral norm
class SNConv2d(nn.Conv2d, SN):
def __init__(self, in_channels, out_channels, kernel_size, stride=1,
padding=0, dilation=1, groups=1, bias=True,
num_svs=1, num_itrs=1, eps=1e-12):
nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, stride,
padding, dilation, groups, bias)
SN.__init__(self, num_svs, num_itrs, out_channels, eps=eps)
def forward(self, x):
return F.conv2d(x, self.W_(), self.bias, self.stride,
self.padding, self.dilation, self.groups)
# Linear layer with spectral norm
class SNLinear(nn.Linear, SN):
def __init__(self, in_features, out_features, bias=True,
num_svs=1, num_itrs=1, eps=1e-12):
nn.Linear.__init__(self, in_features, out_features, bias)
SN.__init__(self, num_svs, num_itrs, out_features, eps=eps)
def forward(self, x):
return F.linear(x, self.W_(), self.bias)
# Embedding layer with spectral norm
# We use num_embeddings as the dim instead of embedding_dim here
# for convenience sake
class SNEmbedding(nn.Embedding, SN):
def __init__(self, num_embeddings, embedding_dim, padding_idx=None,
max_norm=None, norm_type=2, scale_grad_by_freq=False,
sparse=False, _weight=None,
num_svs=1, num_itrs=1, eps=1e-12):
nn.Embedding.__init__(self, num_embeddings, embedding_dim, padding_idx,
max_norm, norm_type, scale_grad_by_freq,
sparse, _weight)
SN.__init__(self, num_svs, num_itrs, num_embeddings, eps=eps)
def forward(self, x):
return F.embedding(x, self.W_())
# A non-local block as used in SA-GAN
# Note that the implementation as described in the paper is largely incorrect;
# refer to the released code for the actual implementation.
class Attention(nn.Module):
def __init__(self, ch, which_conv=SNConv2d, name='attention'):
super(Attention, self).__init__()
# Channel multiplier
self.ch = ch
self.which_conv = which_conv
self.theta = self.which_conv(self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False)
self.phi = self.which_conv(self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False)
self.g = self.which_conv(self.ch, self.ch // 2, kernel_size=1, padding=0, bias=False)
self.o = self.which_conv(self.ch // 2, self.ch, kernel_size=1, padding=0, bias=False)
# Learnable gain parameter
self.gamma = P(torch.tensor(0.), requires_grad=True)
def forward(self, x, y=None):
# Apply convs
theta = self.theta(x)
phi = F.max_pool2d(self.phi(x), [2,2])
g = F.max_pool2d(self.g(x), [2,2])
# Perform reshapes
theta = theta.view(-1, self. ch // 8, x.shape[2] * x.shape[3])
phi = phi.view(-1, self. ch // 8, x.shape[2] * x.shape[3] // 4)
g = g.view(-1, self. ch // 2, x.shape[2] * x.shape[3] // 4)
# Matmul and softmax to get attention maps
beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), -1)
# Attention map times g path
o = self.o(torch.bmm(g, beta.transpose(1,2)).view(-1, self.ch // 2, x.shape[2], x.shape[3]))
return self.gamma * o + x
# Fused batchnorm op
def fused_bn(x, mean, var, gain=None, bias=None, eps=1e-5):
# Apply scale and shift--if gain and bias are provided, fuse them here
# Prepare scale
scale = torch.rsqrt(var + eps)
# If a gain is provided, use it
if gain is not None:
scale = scale * gain
# Prepare shift
shift = mean * scale
# If bias is provided, use it
if bias is not None:
shift = shift - bias
return x * scale - shift
#return ((x - mean) / ((var + eps) ** 0.5)) * gain + bias # The unfused way.
# Manual BN
# Calculate means and variances using mean-of-squares minus mean-squared
def manual_bn(x, gain=None, bias=None, return_mean_var=False, eps=1e-5):
# Cast x to float32 if necessary
float_x = x.float()
# Calculate expected value of x (m) and expected value of x**2 (m2)
# Mean of x
m = torch.mean(float_x, [0, 2, 3], keepdim=True)
# Mean of x squared
m2 = torch.mean(float_x ** 2, [0, 2, 3], keepdim=True)
# Calculate variance as mean of squared minus mean squared.
var = (m2 - m **2)
# Cast back to float 16 if necessary
var = var.type(x.type())
m = m.type(x.type())
# Return mean and variance for updating stored mean/var if requested
if return_mean_var:
return fused_bn(x, m, var, gain, bias, eps), m.squeeze(), var.squeeze()
else:
return fused_bn(x, m, var, gain, bias, eps)
# My batchnorm, supports standing stats
class myBN(nn.Module):
def __init__(self, num_channels, eps=1e-5, momentum=0.1):
super(myBN, self).__init__()
# momentum for updating running stats
self.momentum = momentum
# epsilon to avoid dividing by 0
self.eps = eps
# Momentum
self.momentum = momentum
# Register buffers
self.register_buffer('stored_mean', torch.zeros(num_channels))
self.register_buffer('stored_var', torch.ones(num_channels))
self.register_buffer('accumulation_counter', torch.zeros(1))
# Accumulate running means and vars
self.accumulate_standing = False
# reset standing stats
def reset_stats(self):
self.stored_mean[:] = 0
self.stored_var[:] = 0
self.accumulation_counter[:] = 0
def forward(self, x, gain, bias):
if self.training:
out, mean, var = manual_bn(x, gain, bias, return_mean_var=True, eps=self.eps)
# If accumulating standing stats, increment them
if self.accumulate_standing:
self.stored_mean[:] = self.stored_mean + mean.data
self.stored_var[:] = self.stored_var + var.data
self.accumulation_counter += 1.0
# If not accumulating standing stats, take running averages
else:
self.stored_mean[:] = self.stored_mean * (1 - self.momentum) + mean * self.momentum
self.stored_var[:] = self.stored_var * (1 - self.momentum) + var * self.momentum
return out
# If not in training mode, use the stored statistics
else:
mean = self.stored_mean.view(1, -1, 1, 1)
var = self.stored_var.view(1, -1, 1, 1)
# If using standing stats, divide them by the accumulation counter
if self.accumulate_standing:
mean = mean / self.accumulation_counter
var = var / self.accumulation_counter
return fused_bn(x, mean, var, gain, bias, self.eps)
# Simple function to handle groupnorm norm stylization
def groupnorm(x, norm_style):
# If number of channels specified in norm_style:
if 'ch' in norm_style:
ch = int(norm_style.split('_')[-1])
groups = max(int(x.shape[1]) // ch, 1)
# If number of groups specified in norm style
elif 'grp' in norm_style:
groups = int(norm_style.split('_')[-1])
# If neither, default to groups = 16
else:
groups = 16
return F.group_norm(x, groups)
# Class-conditional bn
# output size is the number of channels, input size is for the linear layers
# Andy's Note: this class feels messy but I'm not really sure how to clean it up
# Suggestions welcome! (By which I mean, refactor this and make a pull request
# if you want to make this more readable/usable).
class ccbn(nn.Module):
def __init__(self, output_size, input_size, which_linear, eps=1e-5, momentum=0.1,
cross_replica=False, mybn=False, norm_style='bn',):
super(ccbn, self).__init__()
self.output_size, self.input_size = output_size, input_size
# Prepare gain and bias layers
self.gain = which_linear(input_size, output_size)
self.bias = which_linear(input_size, output_size)
# epsilon to avoid dividing by 0
self.eps = eps
# Momentum
self.momentum = momentum
# Use cross-replica batchnorm?
self.cross_replica = cross_replica
# Use my batchnorm?
self.mybn = mybn
# Norm style?
self.norm_style = norm_style
if self.cross_replica:
self.bn = SyncBN2d(output_size, eps=self.eps, momentum=self.momentum, affine=False)
elif self.mybn:
self.bn = myBN(output_size, self.eps, self.momentum)
elif self.norm_style in ['bn', 'in']:
self.register_buffer('stored_mean', torch.zeros(output_size))
self.register_buffer('stored_var', torch.ones(output_size))
def forward(self, x, y):
# Calculate class-conditional gains and biases
gain = (1 + self.gain(y)).view(y.size(0), -1, 1, 1)
bias = self.bias(y).view(y.size(0), -1, 1, 1)
# If using my batchnorm
if self.mybn or self.cross_replica:
return self.bn(x, gain=gain, bias=bias)
# else:
else:
if self.norm_style == 'bn':
out = F.batch_norm(x, self.stored_mean, self.stored_var, None, None,
self.training, 0.1, self.eps)
elif self.norm_style == 'in':
out = F.instance_norm(x, self.stored_mean, self.stored_var, None, None,
self.training, 0.1, self.eps)
elif self.norm_style == 'gn':
out = groupnorm(x, self.normstyle)
elif self.norm_style == 'nonorm':
out = x
return out * gain + bias
def extra_repr(self):
s = 'out: {output_size}, in: {input_size},'
s +=' cross_replica={cross_replica}'
return s.format(**self.__dict__)
# Normal, non-class-conditional BN
class bn(nn.Module):
def __init__(self, output_size, eps=1e-5, momentum=0.1,
cross_replica=False, mybn=False):
super(bn, self).__init__()
self.output_size= output_size
# Prepare gain and bias layers
self.gain = P(torch.ones(output_size), requires_grad=True)
self.bias = P(torch.zeros(output_size), requires_grad=True)
# epsilon to avoid dividing by 0
self.eps = eps
# Momentum
self.momentum = momentum
# Use cross-replica batchnorm?
self.cross_replica = cross_replica
# Use my batchnorm?
self.mybn = mybn
if self.cross_replica:
self.bn = SyncBN2d(output_size, eps=self.eps, momentum=self.momentum, affine=False)
elif mybn:
self.bn = myBN(output_size, self.eps, self.momentum)
# Register buffers if neither of the above
else:
self.register_buffer('stored_mean', torch.zeros(output_size))
self.register_buffer('stored_var', torch.ones(output_size))
def forward(self, x, y=None):
if self.cross_replica or self.mybn:
gain = self.gain.view(1,-1,1,1)
bias = self.bias.view(1,-1,1,1)
return self.bn(x, gain=gain, bias=bias)
else:
return F.batch_norm(x, self.stored_mean, self.stored_var, self.gain,
self.bias, self.training, self.momentum, self.eps)
# Generator blocks
# Note that this class assumes the kernel size and padding (and any other
# settings) have been selected in the main generator module and passed in
# through the which_conv arg. Similar rules apply with which_bn (the input
# size [which is actually the number of channels of the conditional info] must
# be preselected)
class GBlock(nn.Module):
def __init__(self, in_channels, out_channels,
which_conv=nn.Conv2d, which_bn=bn, activation=None,
upsample=None):
super(GBlock, self).__init__()
self.in_channels, self.out_channels = in_channels, out_channels
self.which_conv, self.which_bn = which_conv, which_bn
self.activation = activation
self.upsample = upsample
# Conv layers
self.conv1 = self.which_conv(self.in_channels, self.out_channels)
self.conv2 = self.which_conv(self.out_channels, self.out_channels)
self.learnable_sc = in_channels != out_channels or upsample
if self.learnable_sc:
self.conv_sc = self.which_conv(in_channels, out_channels,
kernel_size=1, padding=0)
# Batchnorm layers
self.bn1 = self.which_bn(in_channels)
self.bn2 = self.which_bn(out_channels)
# upsample layers
self.upsample = upsample
def forward(self, x, y):
h = self.activation(self.bn1(x, y))
if self.upsample:
h = self.upsample(h)
x = self.upsample(x)
h = self.conv1(h)
h = self.activation(self.bn2(h, y))
h = self.conv2(h)
if self.learnable_sc:
x = self.conv_sc(x)
return h + x
# Residual block for the discriminator
class DBlock(nn.Module):
def __init__(self, in_channels, out_channels, which_conv=SNConv2d, wide=True,
preactivation=False, activation=None, downsample=None,):
super(DBlock, self).__init__()
self.in_channels, self.out_channels = in_channels, out_channels
# If using wide D (as in SA-GAN and BigGAN), change the channel pattern
self.hidden_channels = self.out_channels if wide else self.in_channels
self.which_conv = which_conv
self.preactivation = preactivation
self.activation = activation
self.downsample = downsample
# Conv layers
self.conv1 = self.which_conv(self.in_channels, self.hidden_channels)
self.conv2 = self.which_conv(self.hidden_channels, self.out_channels)
self.learnable_sc = True if (in_channels != out_channels) or downsample else False
if self.learnable_sc:
self.conv_sc = self.which_conv(in_channels, out_channels,
kernel_size=1, padding=0)
def shortcut(self, x):
if self.preactivation:
if self.learnable_sc:
x = self.conv_sc(x)
if self.downsample:
x = self.downsample(x)
else:
if self.downsample:
x = self.downsample(x)
if self.learnable_sc:
x = self.conv_sc(x)
return x
def forward(self, x):
if self.preactivation:
# h = self.activation(x) # NOT TODAY SATAN
# Andy's note: This line *must* be an out-of-place ReLU or it
# will negatively affect the shortcut connection.
h = F.relu(x)
else:
h = x
h = self.conv1(h)
h = self.conv2(self.activation(h))
if self.downsample:
h = self.downsample(h)
return h + self.shortcut(x)
# dogball
================================================
FILE: BigGAN_utils/logs/BigGAN_ch96_bs256x8.jsonl
================================================
{"itr": 2000, "IS_mean": 2.806771755218506, "IS_std": 0.019480662420392036, "FID": 173.76484159711126, "_stamp": 1551403232.0425167}
{"itr": 4000, "IS_mean": 4.962374687194824, "IS_std": 0.07276841998100281, "FID": 113.86730514283107, "_stamp": 1551422228.743057}
{"itr": 6000, "IS_mean": 6.939817905426025, "IS_std": 0.11417163163423538, "FID": 101.63548498447199, "_stamp": 1551457139.3400874}
{"itr": 8000, "IS_mean": 8.142985343933105, "IS_std": 0.11931543797254562, "FID": 92.0014385772705, "_stamp": 1551476217.2409613}
{"itr": 10000, "IS_mean": 10.355518341064453, "IS_std": 0.09094739705324173, "FID": 83.58068997965364, "_stamp": 1551494854.2419689}
{"itr": 12000, "IS_mean": 11.288347244262695, "IS_std": 0.14952820539474487, "FID": 80.98066299357106, "_stamp": 1551513232.5049698}
{"itr": 14000, "IS_mean": 11.755794525146484, "IS_std": 0.17969024181365967, "FID": 76.80603924280956, "_stamp": 1551531425.150371}
{"itr": 18000, "IS_mean": 13.65534496307373, "IS_std": 0.11151058971881866, "FID": 65.95736694335938, "_stamp": 1551588271.9177916}
{"itr": 20000, "IS_mean": 14.817827224731445, "IS_std": 0.23588882386684418, "FID": 61.32061767578125, "_stamp": 1551606713.6567464}
{"itr": 22000, "IS_mean": 17.16551399230957, "IS_std": 0.19506946206092834, "FID": 53.387969970703125, "_stamp": 1551624876.6513028}
{"itr": 24000, "IS_mean": 19.60654067993164, "IS_std": 0.5591856837272644, "FID": 46.5386962890625, "_stamp": 1551642822.6126688}
{"itr": 26000, "IS_mean": 21.74416732788086, "IS_std": 0.2850531041622162, "FID": 41.595001220703125, "_stamp": 1551663522.6019194}
{"itr": 28000, "IS_mean": 23.923612594604492, "IS_std": 0.41587772965431213, "FID": 37.894744873046875, "_stamp": 1551681794.6567173}
{"itr": 30000, "IS_mean": 25.569377899169922, "IS_std": 0.3333457112312317, "FID": 35.49310302734375, "_stamp": 1551699773.7080302}
{"itr": 32000, "IS_mean": 26.867944717407227, "IS_std": 0.5968036651611328, "FID": 33.4849853515625, "_stamp": 1551717623.887933}
{"itr": 34000, "IS_mean": 28.719074249267578, "IS_std": 0.5698027014732361, "FID": 31.375518798828125, "_stamp": 1551735411.1578612}
{"itr": 36000, "IS_mean": 30.587574005126953, "IS_std": 0.5044271349906921, "FID": 29.432281494140625, "_stamp": 1551783380.6357439}
{"itr": 38000, "IS_mean": 32.08299255371094, "IS_std": 0.49342143535614014, "FID": 28.099456787109375, "_stamp": 1551801179.6495197}
{"itr": 40000, "IS_mean": 34.24657440185547, "IS_std": 0.7709177732467651, "FID": 26.53802490234375, "_stamp": 1551818775.171794}
{"itr": 42000, "IS_mean": 35.891212463378906, "IS_std": 0.7036871314048767, "FID": 25.03021240234375, "_stamp": 1551836329.6873965}
{"itr": 44000, "IS_mean": 38.184898376464844, "IS_std": 0.32996198534965515, "FID": 23.4940185546875, "_stamp": 1551897864.911537}
{"itr": 46000, "IS_mean": 40.239479064941406, "IS_std": 0.7761151194572449, "FID": 22.53167724609375, "_stamp": 1551915406.4840703}
{"itr": 48000, "IS_mean": 41.46656036376953, "IS_std": 1.1031498908996582, "FID": 21.5338134765625, "_stamp": 1551932899.6074848}
{"itr": 50000, "IS_mean": 43.31670379638672, "IS_std": 0.7796809077262878, "FID": 20.53253173828125, "_stamp": 1551950390.345334}
{"itr": 52000, "IS_mean": 45.1517333984375, "IS_std": 1.2925242185592651, "FID": 19.656646728515625, "_stamp": 1551967838.1501615}
{"itr": 54000, "IS_mean": 47.638771057128906, "IS_std": 1.0689665079116821, "FID": 18.898162841796875, "_stamp": 1552044534.5349634}
{"itr": 56000, "IS_mean": 48.87520217895508, "IS_std": 1.1317559480667114, "FID": 18.1248779296875, "_stamp": 1552061763.3080354}
{"itr": 58000, "IS_mean": 49.40987014770508, "IS_std": 1.1866596937179565, "FID": 17.751922607421875, "_stamp": 1552078939.9828825}
{"itr": 60000, "IS_mean": 51.051334381103516, "IS_std": 1.2281248569488525, "FID": 17.19964599609375, "_stamp": 1552096167.889482}
{"itr": 62000, "IS_mean": 52.0235481262207, "IS_std": 0.5391153693199158, "FID": 16.62115478515625, "_stamp": 1552113417.9520617}
{"itr": 64000, "IS_mean": 53.868492126464844, "IS_std": 1.327082633972168, "FID": 16.237335205078125, "_stamp": 1552142961.09602}
{"itr": 66000, "IS_mean": 54.978721618652344, "IS_std": 0.9502049088478088, "FID": 15.81170654296875, "_stamp": 1552162403.2232807}
{"itr": 68000, "IS_mean": 55.73248291015625, "IS_std": 1.0323851108551025, "FID": 15.545623779296875, "_stamp": 1552181112.676657}
{"itr": 70000, "IS_mean": 56.78422927856445, "IS_std": 1.211003303527832, "FID": 15.28369140625, "_stamp": 1552199498.887533}
{"itr": 72000, "IS_mean": 57.972999572753906, "IS_std": 0.8668608665466309, "FID": 14.86395263671875, "_stamp": 1552217782.2738616}
{"itr": 74000, "IS_mean": 58.845054626464844, "IS_std": 1.4297977685928345, "FID": 14.620635986328125, "_stamp": 1552251085.1781816}
{"itr": 76000, "IS_mean": 59.60982131958008, "IS_std": 0.9095696210861206, "FID": 14.360198974609375, "_stamp": 1552270214.9345307}
{"itr": 78000, "IS_mean": 60.71195602416992, "IS_std": 0.960899829864502, "FID": 14.07183837890625, "_stamp": 1552288697.1580262}
{"itr": 80000, "IS_mean": 61.772125244140625, "IS_std": 0.6913255453109741, "FID": 13.781585693359375, "_stamp": 1552307170.0280282}
{"itr": 82000, "IS_mean": 62.98079299926758, "IS_std": 1.4735801219940186, "FID": 13.55389404296875, "_stamp": 1552325252.8553352}
{"itr": 84000, "IS_mean": 64.95240783691406, "IS_std": 0.9018951654434204, "FID": 13.231689453125, "_stamp": 1552344135.3111835}
{"itr": 86000, "IS_mean": 65.13968658447266, "IS_std": 0.8772205114364624, "FID": 13.176849365234375, "_stamp": 1552362429.6782444}
{"itr": 88000, "IS_mean": 65.84476470947266, "IS_std": 1.167534351348877, "FID": 12.87078857421875, "_stamp": 1552380560.7988124}
{"itr": 90000, "IS_mean": 67.41099548339844, "IS_std": 1.6899267435073853, "FID": 12.586517333984375, "_stamp": 1552398550.2060475}
{"itr": 92000, "IS_mean": 68.63685607910156, "IS_std": 1.9431978464126587, "FID": 12.49505615234375, "_stamp": 1552430781.6406457}
{"itr": 94000, "IS_mean": 70.09907531738281, "IS_std": 1.0715738534927368, "FID": 12.047607421875, "_stamp": 1552449001.1950285}
{"itr": 96000, "IS_mean": 70.34623718261719, "IS_std": 1.7962944507598877, "FID": 11.896697998046875, "_stamp": 1552466989.3587568}
{"itr": 98000, "IS_mean": 71.08210754394531, "IS_std": 1.458209753036499, "FID": 11.73046875, "_stamp": 1552484800.7138846}
{"itr": 100000, "IS_mean": 72.24256896972656, "IS_std": 1.3259714841842651, "FID": 11.7386474609375, "_stamp": 1552502538.0269725}
{"itr": 102000, "IS_mean": 73.19488525390625, "IS_std": 1.3439149856567383, "FID": 11.50494384765625, "_stamp": 1552523284.4514356}
{"itr": 104000, "IS_mean": 73.38243103027344, "IS_std": 1.4162707328796387, "FID": 11.374542236328125, "_stamp": 1552541012.0651608}
{"itr": 106000, "IS_mean": 74.95563507080078, "IS_std": 1.089124083518982, "FID": 11.10479736328125, "_stamp": 1552558577.7458107}
{"itr": 108000, "IS_mean": 76.42997741699219, "IS_std": 1.9282453060150146, "FID": 10.998870849609375, "_stamp": 1552576111.9480467}
{"itr": 110000, "IS_mean": 76.89225769042969, "IS_std": 1.4771150350570679, "FID": 10.847015380859375, "_stamp": 1552593659.445132}
{"itr": 112000, "IS_mean": 78.04684448242188, "IS_std": 1.4850096702575684, "FID": 10.772552490234375, "_stamp": 1552616479.5201895}
{"itr": 114000, "IS_mean": 79.67677307128906, "IS_std": 2.0147368907928467, "FID": 10.528045654296875, "_stamp": 1552633850.9315467}
{"itr": 116000, "IS_mean": 79.8828125, "IS_std": 0.978247344493866, "FID": 10.626068115234375, "_stamp": 1552651198.9012825}
{"itr": 118000, "IS_mean": 79.95381164550781, "IS_std": 1.8608143329620361, "FID": 10.46771240234375, "_stamp": 1552668560.4420238}
{"itr": 120000, "IS_mean": 82.37217712402344, "IS_std": 1.8909310102462769, "FID": 10.259033203125, "_stamp": 1552749673.4319007}
{"itr": 122000, "IS_mean": 83.49666595458984, "IS_std": 2.38446044921875, "FID": 9.996185302734375, "_stamp": 1552766698.2706933}
{"itr": 124000, "IS_mean": 83.05189514160156, "IS_std": 1.8844469785690308, "FID": 10.164398193359375, "_stamp": 1552783762.891172}
{"itr": 126000, "IS_mean": 84.27763366699219, "IS_std": 0.9329544901847839, "FID": 10.03509521484375, "_stamp": 1552800953.5724175}
{"itr": 128000, "IS_mean": 85.84852600097656, "IS_std": 2.2698562145233154, "FID": 9.91644287109375, "_stamp": 1552818112.227726}
{"itr": 130000, "IS_mean": 87.356689453125, "IS_std": 2.0958640575408936, "FID": 9.771148681640625, "_stamp": 1552837539.995247}
{"itr": 132000, "IS_mean": 88.72562408447266, "IS_std": 1.7551432847976685, "FID": 9.8258056640625, "_stamp": 1552859685.9305944}
{"itr": 134000, "IS_mean": 88.0631103515625, "IS_std": 1.8199039697647095, "FID": 9.957183837890625, "_stamp": 1552880037.5408435}
{"itr": 136000, "IS_mean": 91.50938415527344, "IS_std": 1.9926033020019531, "FID": 9.876556396484375, "_stamp": 1552899854.652669}
{"itr": 138000, "IS_mean": 93.09217834472656, "IS_std": 2.3062736988067627, "FID": 9.908477783203125, "_stamp": 1552921580.958927}
================================================
FILE: BigGAN_utils/logs/compare_IS.m
================================================
clc
clear all
close all
fclose all;
%% Get All logs and sort them
s = {};
d = dir();
j = 1;
for i = 1:length(d)
if any(strfind(d(i).name,'.jsonl'))
s = [s; d(i).name];
end
end
j = 1;
for i = 1:length(s)
fname = s{i,1};
% Check if the Inception metrics log exists, and if so, plot it
[itr, IS, FID, t] = process_inception_log(fname(1:end - 10), 'log.jsonl');
s{i,2} = itr;
s{i,3} = IS;
s{i,4} = FID;
s{i,5} = max(IS);
s{i,6} = min(FID);
s{i,7} = t;
end
% Sort by Inception Score?
[IS_sorted, IS_index] = sort(cell2mat(s(:,5)));
% Cutoff inception scores below a certain value?
threshold = 22;
IS_index = IS_index(IS_sorted > threshold);
% Sort by FID?
[FID_sorted, FID_index] = sort(cell2mat(s(:,6)));
% Cutoff also based on IS?
% threshold = 0;
FID_index = FID_index(IS_sorted > threshold);
%% Plot things?
cc = hsv(length(IS_index));
legend1 = {};
legend2 = {};
make_axis=true;%false % Turn this on to see the axis out to 1e6 iterations
for i=1:length(IS_index)
legend1 = [legend1; s{IS_index(i), 1}];
figure(1)
plot(s{IS_index(i),2}, s{IS_index(i),3}, 'color', cc(i,:),'linewidth',2)
hold on;
xlabel('itr'); ylabel('IS');
grid on;
if make_axis
axis([0,1e6,0,80]); % 50% grid on;
end
legend(legend1,'Interpreter','none')
%pause(1) % Turn this on to animate stuff
legend2 = [legend2; s{IS_index(i), 1}];
figure(2)
plot(s{IS_index(i),2}, s{IS_index(i),4}, 'color', cc(i,:),'linewidth',2)
hold on;
xlabel('itr'); ylabel('FID');
j = j + 1;
grid on;
if make_axis
axis([0,1e6,0,50]);% grid on;
end
legend(legend2, 'Interpreter','none')
end
%% Quick script to plot IS versus timesteps
if 0
figure(3);
this_index=4;
subplot(2,1,1);
%plot(s{this_index, 2}(2:end), s{this_index, 7}(2:end) - s{this_index, 7}(1:end-1), 'r*');
% xlabel('Iteration');ylabel('\Delta T')
plot(s{this_index, 2}, s{this_index, 7}, 'r*');
xlabel('Iteration');ylabel('T')
subplot(2,1,2);
plot(s{this_index, 2}, s{this_index, 3}, 'r', 'linewidth',2);
xlabel('Iteration'), ylabel('Inception score')
title(s{this_index,1})
end
================================================
FILE: BigGAN_utils/logs/metalog.txt
================================================
datetime: 2019-03-18 13:27:59.181225
config: {'dataset': 'I128_hdf5', 'augment': False, 'num_workers': 8, 'pin_memory': True, 'shuffle': True, 'load_in_mem': True, 'use_multiepoch_sampler': True, 'model': 'model', 'G_param': 'SN', 'D_param': 'SN', 'G_ch': 96, 'D_ch': 96, 'G_depth': 1, 'D_depth': 1, 'D_wide': True, 'G_shared': True, 'shared_dim': 128, 'dim_z': 120, 'z_var': 1.0, 'hier': True, 'cross_replica': False, 'mybn': False, 'G_nl': 'inplace_relu', 'D_nl': 'inplace_relu', 'G_attn': '64', 'D_attn': '64', 'norm_style': 'bn', 'seed': 0, 'G_init': 'ortho', 'D_init': 'ortho', 'skip_init': True, 'G_lr': 0.0001, 'D_lr': 0.0004, 'G_B1': 0.0, 'D_B1': 0.0, 'G_B2': 0.999, 'D_B2': 0.999, 'batch_size': 256, 'G_batch_size': 0, 'num_G_accumulations': 8, 'num_D_steps': 1, 'num_D_accumulations': 8, 'split_D': False, 'num_epochs': 400, 'parallel': True, 'G_fp16': False, 'D_fp16': False, 'D_mixed_precision': False, 'G_mixed_precision': False, 'accumulate_stats': False, 'num_standing_accumulations': 16, 'G_eval_mode': True, 'save_every': 500, 'num_save_copies': 2, 'num_best_copies': 5, 'which_best': 'IS', 'no_fid': False, 'test_every': 2000, 'num_inception_images': 50000, 'hashname': False, 'base_root': '', 'dataset_root': 'data', 'weights_root': 'weights', 'logs_root': 'logs', 'samples_root': 'samples', 'pbar': 'mine', 'name_suffix': '', 'experiment_name': 'Jade_BigGAN_B1_bs256x8_fp32', 'config_from_name': False, 'ema': True, 'ema_decay': 0.9999, 'use_ema': True, 'ema_start': 20000, 'adam_eps': 1e-06, 'BN_eps': 1e-05, 'SN_eps': 1e-06, 'num_G_SVs': 1, 'num_D_SVs': 1, 'num_G_SV_itrs': 1, 'num_D_SV_itrs': 1, 'G_ortho': 0.0, 'D_ortho': 0.0, 'toggle_grads': True, 'which_train_fn': 'GAN', 'load_weights': '', 'resume': True, 'logstyle': '%3.3e', 'log_G_spectra': False, 'log_D_spectra': False, 'sv_log_interval': 10, 'resolution': 128, 'n_classes': 1000, 'G_activation': ReLU(inplace), 'D_activation': ReLU(inplace)}
state: {'itr': 137500, 'epoch': 2, 'save_num': 0, 'save_best_num': 1, 'best_IS': 91.509384, 'best_FID': tensor(9.7711, 'config': {'dataset': 'I128_hdf5', 'augment': False, 'num_workers': 8, 'pin_memory': True, 'shuffle': True, 'load_in_mem': True, 'use_multiepoch_sampler': True, 'model': 'model', 'G_param': 'SN', 'D_param': 'SN', 'G_ch': 96, 'D_ch': 96, 'D_wide': True, 'G_shared': True, 'shared_dim': 128, 'dim_z': 120, 'hier': True, 'cross_replica': False, 'mybn': False, 'G_nl': 'inplace_relu', 'D_nl': 'inplace_relu', 'G_attn': '64', 'D_attn': '64', 'norm_style': 'bn', 'seed': 0, 'G_init': 'ortho', 'D_init': 'ortho', 'skip_init': False, 'G_lr': 0.0001, 'D_lr': 0.0004, 'G_B1': 0.0, 'D_B1': 0.0, 'G_B2': 0.999, 'D_B2': 0.999, 'batch_size': 256, 'G_batch_size': 0, 'num_G_accumulations': 8, 'num_D_steps': 1, 'num_D_accumulations': 8, 'split_D': False, 'num_epochs': 100, 'parallel': True, 'G_fp16': False, 'D_fp16': False, 'D_mixed_precision': False, 'G_mixed_precision': False, 'accumulate_stats': False, 'num_standing_accumulations': 16, 'BN_sync': False, 'G_eval_mode': True, 'save_every': 500, 'num_save_copies': 2, 'num_best_copies': 5, 'which_best': 'IS', 'no_fid': False, 'test_every': 2000, 'num_inception_images': 50000, 'hashname': False, 'base_root': '', 'dataset_root': 'data', 'weights_root': 'weights', 'logs_root': 'logs', 'samples_root': 'samples', 'pbar': 'mine', 'name_suffix': '', 'experiment_name': 'Jade_BigGAN_B1_bs256x8_fp32', 'ema': True, 'ema_decay': 0.9999, 'use_ema': True, 'ema_start': 20000, 'adam_eps': 1e-06, 'BN_eps': 1e-05, 'SN_eps': 1e-06, 'num_G_SVs': 1, 'num_D_SVs': 1, 'num_G_SV_itrs': 1, 'num_D_SV_itrs': 1, 'G_ortho': 0.0, 'D_ortho': 0.0, 'toggle_grads': True, 'which_train_fn': 'GAN', 'load_weights': '', 'resume': False, 'logstyle': '%3.3e', 'log_G_spectra': False, 'log_D_spectra': False, 'sv_log_interval': 10, 'resolution': 128, 'n_classes': 1000, 'G_activation': ReLU(inplace), 'D_activation': ReLU(inplace)}}
================================================
FILE: BigGAN_utils/logs/process_inception_log.m
================================================
function [itr, IS, FID, t] = process_inception_log(fname, which_log)
f = sprintf('%s_%s',fname, which_log);%'G_loss.log');
fid = fopen(f,'r');
itr = [];
IS = [];
FID = [];
t = [];
i = 1;
while ~feof(fid);
s = fgets(fid);
parsed = sscanf(s,'{"itr": %d, "IS_mean": %f, "IS_std": %f, "FID": %f, "_stamp": %f}');
itr(i) = parsed(1);
IS(i) = parsed(2);
FID(i) = parsed(4);
t(i) = parsed(5);
i = i + 1;
end
fclose(fid);
end
================================================
FILE: BigGAN_utils/logs/process_training.m
================================================
clc
clear all
close all
fclose all;
%% Get all training logs for a given run
target_dir = '.';
s = {};
nm = {};
d = dir(target_dir);
j = 1;
for i = 1:length(d)
if any(strfind(d(i).name,'.log'))
s = [s; sprintf('%s\\%s', target_dir, d(i).name)];
nm = [nm; d(i).name];
end
end
%% Loop over training logs and acquire data
D_count = 0;
G_count = 0;
for i = 1:length(s)
fname = s{i,1};
fid = fopen(s{i,1},'r');
% Prepare bookkeeping for sv0
if any(strfind(s{i,1},'sv'))
if any(strfind(s{i,1},'G_'))
G_count = G_count +1;
else
D_count = D_count + 1;
end
end
itr = [];
val = [];
j = 1;
while ~feof(fid);
line = fgets(fid);
parsed = sscanf(line, '%d: %e');
itr(j) = parsed(1);
val(j) = parsed(2);
j = j + 1;
end
s{i,2} = itr;
s{i,3} = val;
fclose(fid);
end
%% Plot SVs and losses
close all;
Gcc = hsv(G_count);
Dcc = hsv(D_count);
gi = 1;
di = 1;
li = 1;
legendG = {};
legendD = {};
legendL = {};
thresh=2; % wavelet denoising threshold
losses = {};
for i=1:length(s)
if any(strfind(s{i,1},'D_loss_real.log')) || any(strfind(s{i,1},'D_loss_fake.log')) || any(strfind(s{i,1},'G_loss.log'))
% Select colors
if any(strfind(s{i,1},'D_loss_real.log'))
color1 = [0.7,0.7,1.0];
color2 = [0, 0, 1];
dlr = {s{i,2}, s{i,3}, wden(s{i,3},'sqtwolog','s','mln', thresh, 'sym4'), color1, color2};
losses = [losses; dlr];
elseif any(strfind(s{i,1},'D_loss_fake.log'))
color1 = [0.7,1.0,0.7];
color2 = [0, 1, 0];
dlf = {s{i,2},s{i,3} wden(s{i,3},'sqtwolog','s','mln', thresh, 'sym4'), color1, color2};
losses = [losses; dlf];
else % g loss
color1 = [1.0, 0.7,0.7];
color2 = [1, 0, 0];
gl = {s{i,2},s{i,3}, wden(s{i,3},'sqtwolog','s','mln', thresh, 'sym4'), color1 color2};
losses = [losses; gl];
end
figure(1); hold on;
% Plot the unsmoothed losses; we'll plot the smoothed losses later
plot(s{i,2},s{i,3},'color', color1, 'HandleVisibility','off');
legendL = [legendL; nm{i}];
continue
end
if any(strfind(s{i,1},'G_'))
legendG = [legendG; nm{i}];
figure(2); hold on;
plot(s{i,2},s{i,3},'color',Gcc(gi,:),'linewidth',2);
gi = gi+1;
elseif any(strfind(s{i,1},'D_'))
legendD = [legendD; nm{i}];
figure(3); hold on;
plot(s{i,2},s{i,3},'color',Dcc(di,:),'linewidth',2);
di = di+1;
else
s{i,1} % Debug print to show the name of the log that was not processed.
end
end
figure(1);
% Plot the smoothed losses last
for i = 1:3
% plot(losses{i,1}, losses{i,2},'color', losses{i,4}, 'HandleVisibility','off');
plot(losses{i,1},losses{i,3},'color',losses{i,5});
end
legend(legendL, 'Interpreter', 'none'); title('Losses'); xlabel('Generator itr'); ylabel('loss'); axis([0, max(s{end,2}), -1, 4]);
figure(2); legend(legendG,'Interpreter','none'); title('Singular Values in G'); xlabel('Generator itr'); ylabel('SV0');
figure(3); legend(legendD, 'Interpreter', 'none'); title('Singular Values in D'); xlabel('Generator itr'); ylabel('SV0');
================================================
FILE: BigGAN_utils/losses.py
================================================
import torch
import torch.nn.functional as F
# DCGAN loss
def loss_dcgan_dis(dis_fake, dis_real):
L1 = torch.mean(F.softplus(-dis_real))
L2 = torch.mean(F.softplus(dis_fake))
return L1, L2
def loss_dcgan_gen(dis_fake):
loss = torch.mean(F.softplus(-dis_fake))
return loss
# Hinge Loss
def loss_hinge_dis(dis_fake, dis_real):
loss_real = torch.mean(F.relu(1. - dis_real))
loss_fake = torch.mean(F.relu(1. + dis_fake))
return loss_real, loss_fake
# def loss_hinge_dis(dis_fake, dis_real): # This version returns a single loss
# loss = torch.mean(F.relu(1. - dis_real))
# loss += torch.mean(F.relu(1. + dis_fake))
# return loss
def loss_hinge_gen(dis_fake):
loss = -torch.mean(dis_fake)
return loss
# Default to hinge loss
generator_loss = loss_hinge_gen
discriminator_loss = loss_hinge_dis
================================================
FILE: BigGAN_utils/make_hdf5.py
================================================
""" Convert dataset to HDF5
This script preprocesses a dataset and saves it (images and labels) to
an HDF5 file for improved I/O. """
import os
import sys
from argparse import ArgumentParser
from tqdm import tqdm, trange
import h5py as h5
import numpy as np
import torch
import torchvision.datasets as dset
import torchvision.transforms as transforms
from torchvision.utils import save_image
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import utils
def prepare_parser():
usage = 'Parser for ImageNet HDF5 scripts.'
parser = ArgumentParser(description=usage)
parser.add_argument(
'--dataset', type=str, default='I128',
help='Which Dataset to train on, out of I128, I256, C10, C100;'
'Append "_hdf5" to use the hdf5 version for ISLVRC (default: %(default)s)')
parser.add_argument(
'--data_root', type=str, default='data',
help='Default location where data is stored (default: %(default)s)')
parser.add_argument(
'--batch_size', type=int, default=256,
help='Default overall batchsize (default: %(default)s)')
parser.add_argument(
'--num_workers', type=int, default=16,
help='Number of dataloader workers (default: %(default)s)')
parser.add_argument(
'--chunk_size', type=int, default=500,
help='Default overall batchsize (default: %(default)s)')
parser.add_argument(
'--compression', action='store_true', default=False,
help='Use LZF compression? (default: %(default)s)')
return parser
def run(config):
if 'hdf5' in config['dataset']:
raise ValueError('Reading from an HDF5 file which you will probably be '
'about to overwrite! Override this error only if you know '
'what you''re doing!')
# Get image size
config['image_size'] = utils.imsize_dict[config['dataset']]
# Update compression entry
config['compression'] = 'lzf' if config['compression'] else None #No compression; can also use 'lzf'
# Get dataset
kwargs = {'num_workers': config['num_workers'], 'pin_memory': False, 'drop_last': False}
train_loader = utils.get_data_loaders(dataset=config['dataset'],
batch_size=config['batch_size'],
shuffle=False,
data_root=config['data_root'],
use_multiepoch_sampler=False,
**kwargs)[0]
# HDF5 supports chunking and compression. You may want to experiment
# with different chunk sizes to see how it runs on your machines.
# Chunk Size/compression Read speed @ 256x256 Read speed @ 128x128 Filesize @ 128x128 Time to write @128x128
# 1 / None 20/s
# 500 / None ramps up to 77/s 102/s 61GB 23min
# 500 / LZF 8/s 56GB 23min
# 1000 / None 78/s
# 5000 / None 81/s
# auto:(125,1,16,32) / None 11/s 61GB
print('Starting to load %s into an HDF5 file with chunk size %i and compression %s...' % (config['dataset'], config['chunk_size'], config['compression']))
# Loop over train loader
for i,(x,y) in enumerate(tqdm(train_loader)):
# Stick X into the range [0, 255] since it's coming from the train loader
x = (255 * ((x + 1) / 2.0)).byte().numpy()
# Numpyify y
y = y.numpy()
# If we're on the first batch, prepare the hdf5
if i==0:
with h5.File(config['data_root'] + '/ILSVRC%i.hdf5' % config['image_size'], 'w') as f:
print('Producing dataset of len %d' % len(train_loader.dataset))
imgs_dset = f.create_dataset('imgs', x.shape,dtype='uint8', maxshape=(len(train_loader.dataset), 3, config['image_size'], config['image_size']),
chunks=(config['chunk_size'], 3, config['image_size'], config['image_size']), compression=config['compression'])
print('Image chunks chosen as ' + str(imgs_dset.chunks))
imgs_dset[...] = x
labels_dset = f.create_dataset('labels', y.shape, dtype='int64', maxshape=(len(train_loader.dataset),), chunks=(config['chunk_size'],), compression=config['compression'])
print('Label chunks chosen as ' + str(labels_dset.chunks))
labels_dset[...] = y
# Else append to the hdf5
else:
with h5.File(config['data_root'] + '/ILSVRC%i.hdf5' % config['image_size'], 'a') as f:
f['imgs'].resize(f['imgs'].shape[0] + x.shape[0], axis=0)
f['imgs'][-x.shape[0]:] = x
f['labels'].resize(f['labels'].shape[0] + y.shape[0], axis=0)
f['labels'][-y.shape[0]:] = y
def main():
# parse command line and run
parser = prepare_parser()
config = vars(parser.parse_args())
print(config)
run(config)
if __name__ == '__main__':
main()
================================================
FILE: BigGAN_utils/sample.py
================================================
''' Sample
This script loads a pretrained net and a weightsfile and sample '''
import functools
import math
import numpy as np
from tqdm import tqdm, trange
import torch
import torch.nn as nn
from torch.nn import init
import torch.optim as optim
import torch.nn.functional as F
from torch.nn import Parameter as P
import torchvision
# Import my stuff
import inception_utils
import utils
import losses
def run(config):
# Prepare state dict, which holds things like epoch # and itr #
state_dict = {'itr': 0, 'epoch': 0, 'save_num': 0, 'save_best_num': 0,
'best_IS': 0, 'best_FID': 999999, 'config': config}
# Optionally, get the configuration from the state dict. This allows for
# recovery of the config provided only a state dict and experiment name,
# and can be convenient for writing less verbose sample shell scripts.
if config['config_from_name']:
utils.load_weights(None, None, state_dict, config['weights_root'],
config['experiment_name'], config['load_weights'], None,
strict=False, load_optim=False)
# Ignore items which we might want to overwrite from the command line
for item in state_dict['config']:
if item not in ['z_var', 'base_root', 'batch_size', 'G_batch_size', 'use_ema', 'G_eval_mode']:
config[item] = state_dict['config'][item]
# update config (see train.py for explanation)
config['resolution'] = utils.imsize_dict[config['dataset']]
config['n_classes'] = utils.nclass_dict[config['dataset']]
config['G_activation'] = utils.activation_dict[config['G_nl']]
config['D_activation'] = utils.activation_dict[config['D_nl']]
config = utils.update_config_roots(config)
config['skip_init'] = True
config['no_optim'] = True
device = 'cuda'
# Seed RNG
utils.seed_rng(config['seed'])
# Setup cudnn.benchmark for free speed
torch.backends.cudnn.benchmark = True
# Import the model--this line allows us to dynamically select different files.
model = __import__(config['model'])
experiment_name = (config['experiment_name'] if config['experiment_name']
else utils.name_from_config(config))
print('Experiment name is %s' % experiment_name)
G = model.Generator(**config).cuda()
utils.count_parameters(G)
# Load weights
print('Loading weights...')
# Here is where we deal with the ema--load ema weights or load normal weights
utils.load_weights(G if not (config['use_ema']) else None, None, state_dict,
config['weights_root'], experiment_name, config['load_weights'],
G if config['ema'] and config['use_ema'] else None,
strict=False, load_optim=False)
# Update batch size setting used for G
G_batch_size = max(config['G_batch_size'], config['batch_size'])
z_, y_ = utils.prepare_z_y(G_batch_size, G.dim_z, config['n_classes'],
device=device, fp16=config['G_fp16'],
z_var=config['z_var'])
if config['G_eval_mode']:
print('Putting G in eval mode..')
G.eval()
else:
print('G is in %s mode...' % ('training' if G.training else 'eval'))
#Sample function
sample = functools.partial(utils.sample, G=G, z_=z_, y_=y_, config=config)
if config['accumulate_stats']:
print('Accumulating standing stats across %d accumulations...' % config['num_standing_accumulations'])
utils.accumulate_standing_stats(G, z_, y_, config['n_classes'],
config['num_standing_accumulations'])
# Sample a number of images and save them to an NPZ, for use with TF-Inception
if config['sample_npz']:
# Lists to hold images and labels for images
x, y = [], []
print('Sampling %d images and saving them to npz...' % config['sample_num_npz'])
for i in trange(int(np.ceil(config['sample_num_npz'] / float(G_batch_size)))):
with torch.no_grad():
images, labels = sample()
x += [np.uint8(255 * (images.cpu().numpy() + 1) / 2.)]
y += [labels.cpu().numpy()]
x = np.concatenate(x, 0)[:config['sample_num_npz']]
y = np.concatenate(y, 0)[:config['sample_num_npz']]
print('Images shape: %s, Labels shape: %s' % (x.shape, y.shape))
npz_filename = '%s/%s/samples.npz' % (config['samples_root'], experiment_name)
print('Saving npz to %s...' % npz_filename)
np.savez(npz_filename, **{'x' : x, 'y' : y})
# Prepare sample sheets
if config['sample_sheets']:
print('Preparing conditional sample sheets...')
utils.sample_sheet(G, classes_per_sheet=utils.classes_per_sheet_dict[config['dataset']],
num_classes=config['n_classes'],
samples_per_class=10, parallel=config['parallel'],
samples_root=config['samples_root'],
experiment_name=experiment_name,
folder_number=config['sample_sheet_folder_num'],
z_=z_,)
# Sample interp sheets
if config['sample_interps']:
print('Preparing interp sheets...')
for fix_z, fix_y in zip([False, False, True], [False, True, False]):
utils.interp_sheet(G, num_per_sheet=16, num_midpoints=8,
num_classes=config['n_classes'],
parallel=config['parallel'],
samples_root=config['samples_root'],
experiment_name=experiment_name,
folder_number=config['sample_sheet_folder_num'],
sheet_number=0,
fix_z=fix_z, fix_y=fix_y, device='cuda')
# Sample random sheet
if config['sample_random']:
print('Preparing random sample sheet...')
images, labels = sample()
torchvision.utils.save_image(images.float(),
'%s/%s/random_samples.jpg' % (config['samples_root'], experiment_name),
nrow=int(G_batch_size**0.5),
normalize=True)
# Get Inception Score and FID
get_inception_metrics = inception_utils.prepare_inception_metrics(config['dataset'], config['parallel'], config['no_fid'])
# Prepare a simple function get metrics that we use for trunc curves
def get_metrics():
sample = functools.partial(utils.sample, G=G, z_=z_, y_=y_, config=config)
IS_mean, IS_std, FID = get_inception_metrics(sample, config['num_inception_images'], num_splits=10, prints=False)
# Prepare output string
outstring = 'Using %s weights ' % ('ema' if config['use_ema'] else 'non-ema')
outstring += 'in %s mode, ' % ('eval' if config['G_eval_mode'] else 'training')
outstring += 'with noise variance %3.3f, ' % z_.var
outstring += 'over %d images, ' % config['num_inception_images']
if config['accumulate_stats'] or not config['G_eval_mode']:
outstring += 'with batch size %d, ' % G_batch_size
if config['accumulate_stats']:
outstring += 'using %d standing stat accumulations, ' % config['num_standing_accumulations']
outstring += 'Itr %d: PYTORCH UNOFFICIAL Inception Score is %3.3f +/- %3.3f, PYTORCH UNOFFICIAL FID is %5.4f' % (state_dict['itr'], IS_mean, IS_std, FID)
print(outstring)
if config['sample_inception_metrics']:
print('Calculating Inception metrics...')
get_metrics()
# Sample truncation curve stuff. This is basically the same as the inception metrics code
if config['sample_trunc_curves']:
start, step, end = [float(item) for item in config['sample_trunc_curves'].split('_')]
print('Getting truncation values for variance in range (%3.3f:%3.3f:%3.3f)...' % (start, step, end))
for var in np.arange(start, end + step, step):
z_.var = var
# Optionally comment this out if you want to run with standing stats
# accumulated at one z variance setting
if config['accumulate_stats']:
utils.accumulate_standing_stats(G, z_, y_, config['n_classes'],
config['num_standing_accumulations'])
get_metrics()
def main():
# parse command line and run
parser = utils.prepare_parser()
parser = utils.add_sample_parser(parser)
config = vars(parser.parse_args())
print(config)
run(config)
if __name__ == '__main__':
main()
================================================
FILE: BigGAN_utils/scripts/launch_BigGAN_bs256x8.sh
================================================
#!/bin/bash
python train.py \
--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 256 --load_in_mem \
--num_G_accumulations 8 --num_D_accumulations 8 \
--num_D_steps 1 --G_lr 1e-4 --D_lr 4e-4 --D_B2 0.999 --G_B2 0.999 \
--G_attn 64 --D_attn 64 \
--G_nl inplace_relu --D_nl inplace_relu \
--SN_eps 1e-6 --BN_eps 1e-5 --adam_eps 1e-6 \
--G_ortho 0.0 \
--G_shared \
--G_init ortho --D_init ortho \
--hier --dim_z 120 --shared_dim 128 \
--G_eval_mode \
--G_ch 96 --D_ch 96 \
--ema --use_ema --ema_start 20000 \
--test_every 2000 --save_every 1000 --num_best_copies 5 --num_save_copies 2 --seed 0 \
--use_multiepoch_sampler \
================================================
FILE: BigGAN_utils/scripts/launch_BigGAN_bs512x4.sh
================================================
#!/bin/bash
python train.py \
--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 512 --load_in_mem \
--num_G_accumulations 4 --num_D_accumulations 4 \
--num_D_steps 1 --G_lr 1e-4 --D_lr 4e-4 --D_B2 0.999 --G_B2 0.999 \
--G_attn 64 --D_attn 64 \
--G_nl inplace_relu --D_nl inplace_relu \
--SN_eps 1e-6 --BN_eps 1e-5 --adam_eps 1e-6 \
--G_ortho 0.0 \
--G_shared \
--G_init ortho --D_init ortho \
--hier --dim_z 120 --shared_dim 128 \
--G_eval_mode \
--G_ch 96 --D_ch 96 \
--ema --use_ema --ema_start 20000 \
--test_every 2000 --save_every 1000 --num_best_copies 5 --num_save_copies 2 --seed 0 \
--use_multiepoch_sampler \
================================================
FILE: BigGAN_utils/scripts/launch_BigGAN_ch64_bs256x8.sh
================================================
#!/bin/bash
python train.py \
--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 256 --load_in_mem \
--num_G_accumulations 8 --num_D_accumulations 8 \
--num_D_steps 1 --G_lr 1e-4 --D_lr 4e-4 --D_B2 0.999 --G_B2 0.999 \
--G_attn 64 --D_attn 64 \
--G_nl inplace_relu --D_nl inplace_relu \
--SN_eps 1e-6 --BN_eps 1e-5 --adam_eps 1e-6 \
--G_ortho 0.0 \
--G_shared \
--G_init ortho --D_init ortho \
--hier --dim_z 120 --shared_dim 128 \
--G_eval_mode \
--G_ch 64 --G_ch 64 \
--ema --use_ema --ema_start 20000 \
--test_every 2000 --save_every 1000 --num_best_copies 5 --num_save_copies 2 --seed 0 \
--use_multiepoch_sampler
================================================
FILE: BigGAN_utils/scripts/launch_BigGAN_deep.sh
================================================
#!/bin/bash
python train.py \
--model BigGANdeep \
--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 256 \
--num_G_accumulations 8 --num_D_accumulations 8 \
--num_D_steps 1 --G_lr 1e-4 --D_lr 4e-4 --D_B2 0.999 --G_B2 0.999 \
--G_attn 64 --D_attn 64 \
--G_ch 128 --D_ch 128 \
--G_depth 2 --D_depth 2 \
--G_nl inplace_relu --D_nl inplace_relu \
--SN_eps 1e-6 --BN_eps 1e-5 --adam_eps 1e-6 \
--G_ortho 0.0 \
--G_shared \
--G_init ortho --D_init ortho \
--hier --dim_z 128 --shared_dim 128 \
--ema --use_ema --ema_start 20000 --G_eval_mode \
--test_every 2000 --save_every 500 --num_best_copies 5 --num_save_copies 2 --seed 0 \
--use_multiepoch_sampler \
================================================
FILE: BigGAN_utils/scripts/launch_SAGAN_bs128x2_ema.sh
================================================
#!/bin/bash
python train.py \
--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 128 \
--num_G_accumulations 2 --num_D_accumulations 2 \
--num_D_steps 1 --G_lr 1e-4 --D_lr 4e-4 --D_B2 0.999 --G_B2 0.999 \
--G_attn 64 --D_attn 64 \
--G_nl relu --D_nl relu \
--SN_eps 1e-8 --BN_eps 1e-5 --adam_eps 1e-8 \
--G_ortho 0.0 \
--G_init xavier --D_init xavier \
--ema --use_ema --ema_start 2000 --G_eval_mode \
--test_every 2000 --save_every 1000 --num_best_copies 5 --num_save_copies 2 --seed 0 \
--name_suffix SAGAN_ema \
================================================
FILE: BigGAN_utils/scripts/launch_SNGAN.sh
================================================
#!/bin/bash
python train.py \
--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 64 \
--num_G_accumulations 1 --num_D_accumulations 1 \
--num_D_steps 5 --G_lr 2e-4 --D_lr 2e-4 --D_B2 0.900 --G_B2 0.900 \
--G_attn 0 --D_attn 0 \
--G_nl relu --D_nl relu \
--SN_eps 1e-8 --BN_eps 1e-5 --adam_eps 1e-8 \
--G_ortho 0.0 \
--D_thin \
--G_init xavier --D_init xavier \
--G_eval_mode \
--test_every 2000 --save_every 1000 --num_best_copies 5 --num_save_copies 2 --seed 0 \
--name_suffix SNGAN \
================================================
FILE: BigGAN_utils/scripts/launch_cifar_ema.sh
================================================
#!/bin/bash
CUDA_VISIBLE_DEVICES=0,1 python train.py \
--shuffle --batch_size 50 --parallel \
--num_G_accumulations 1 --num_D_accumulations 1 --num_epochs 500 \
--num_D_steps 4 --G_lr 2e-4 --D_lr 2e-4 \
--dataset C10 \
--G_ortho 0.0 \
--G_attn 0 --D_attn 0 \
--G_init N02 --D_init N02 \
--ema --use_ema --ema_start 1000 \
--test_every 5000 --save_every 2000 --num_best_copies 5 --num_save_copies 2 --seed 0
================================================
FILE: BigGAN_utils/scripts/sample_BigGAN_bs256x8.sh
================================================
# use z_var to change the variance of z for all the sampling
# use --mybn --accumulate_stats --num_standing_accumulations 32 to
# use running stats
python sample.py \
--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 256 \
--num_G_accumulations 8 --num_D_accumulations 8 \
--num_D_steps 1 --G_lr 1e-4 --D_lr 4e-4 --D_B2 0.999 --G_B2 0.999 \
--G_attn 64 --D_attn 64 \
--G_ch 96 --D_ch 96 \
--G_nl inplace_relu --D_nl inplace_relu \
--SN_eps 1e-6 --BN_eps 1e-5 --adam_eps 1e-6 \
--G_ortho 0.0 \
--G_shared \
--G_init ortho --D_init ortho --skip_init \
--hier --dim_z 120 --shared_dim 128 \
--ema --ema_start 20000 \
--use_multiepoch_sampler \
--test_every 2000 --save_every 1000 --num_best_copies 5 --num_save_copies 2 --seed 0 \
--skip_init --G_batch_size 512 --use_ema --G_eval_mode --sample_trunc_curves 0.05_0.05_1.0 \
--sample_inception_metrics --sample_npz --sample_random --sample_sheets --sample_interps
================================================
FILE: BigGAN_utils/scripts/sample_cifar_ema.sh
================================================
#!/bin/bash
CUDA_VISIBLE_DEVICES=0,1 python sample.py \
--shuffle --batch_size 50 --G_batch_size 256 --parallel \
--num_G_accumulations 1 --num_D_accumulations 1 --num_epochs 500 \
--num_D_steps 4 --G_lr 2e-4 --D_lr 2e-4 \
--dataset C10 \
--G_ortho 0.0 \
--G_attn 0 --D_attn 0 \
--G_init N02 --D_init N02 \
--ema --use_ema --ema_start 1000 \
--test_ever
gitextract_02nfnavz/ ├── BigGAN_utils/ │ ├── BigGAN.py │ ├── BigGANdeep.py │ ├── LICENSE │ ├── README.md │ ├── TFHub/ │ │ ├── biggan_v1.py │ │ └── converter.py │ ├── __init__.py │ ├── animal_hash.py │ ├── binary_utils.py │ ├── calculate_inception_moments.py │ ├── datasets.py │ ├── inception_tf13.py │ ├── inception_utils.py │ ├── layers.py │ ├── logs/ │ │ ├── BigGAN_ch96_bs256x8.jsonl │ │ ├── compare_IS.m │ │ ├── metalog.txt │ │ ├── process_inception_log.m │ │ └── process_training.m │ ├── losses.py │ ├── make_hdf5.py │ ├── sample.py │ ├── scripts/ │ │ ├── launch_BigGAN_bs256x8.sh │ │ ├── launch_BigGAN_bs512x4.sh │ │ ├── launch_BigGAN_ch64_bs256x8.sh │ │ ├── launch_BigGAN_deep.sh │ │ ├── launch_SAGAN_bs128x2_ema.sh │ │ ├── launch_SNGAN.sh │ │ ├── launch_cifar_ema.sh │ │ ├── sample_BigGAN_bs256x8.sh │ │ ├── sample_cifar_ema.sh │ │ └── utils/ │ │ ├── duplicate.sh │ │ └── prepare_data.sh │ ├── sync_batchnorm/ │ │ ├── __init__.py │ │ ├── batchnorm.py │ │ ├── batchnorm_reimpl.py │ │ ├── comm.py │ │ ├── replicate.py │ │ └── unittest.py │ ├── train.py │ ├── train_fns.py │ ├── utils.py │ └── weights/ │ └── README.md ├── DiffAugment_pytorch.py ├── LICENSE ├── README.md ├── fusedream_generator.py └── fusedream_utils.py
SYMBOL INDEX (296 symbols across 23 files)
FILE: BigGAN_utils/BigGAN.py
function G_arch (line 19) | def G_arch(ch=64, attention='64', ksize='333333', dilation='111111'):
class Generator (line 54) | class Generator(nn.Module):
method __init__ (line 55) | def __init__(self, G_ch=64, dim_z=128, bottom_width=4, resolution=128,
method init_weights (line 210) | def init_weights(self):
method forward (line 231) | def forward(self, z, y, w_y=None):
method forward_org (line 264) | def forward_org(self, z, y):
function D_arch (line 289) | def D_arch(ch=64, attention='64',ksize='333333', dilation='111111'):
class Discriminator (line 317) | class Discriminator(nn.Module):
method __init__ (line 319) | def __init__(self, D_ch=64, D_wide=True, resolution=128,
method init_weights (line 408) | def init_weights(self):
method forward (line 425) | def forward(self, x, y=None):
class G_D (line 442) | class G_D(nn.Module):
method __init__ (line 443) | def __init__(self, G, D):
method forward (line 448) | def forward(self, z, gy, x=None, dy=None, train_G=False, return_G_z=Fa...
FILE: BigGAN_utils/BigGANdeep.py
class GBlock (line 23) | class GBlock(nn.Module):
method __init__ (line 24) | def __init__(self, in_channels, out_channels,
method forward (line 48) | def forward(self, x, y):
function G_arch (line 67) | def G_arch(ch=64, attention='64', ksize='333333', dilation='111111'):
class Generator (line 96) | class Generator(nn.Module):
method __init__ (line 97) | def __init__(self, G_ch=64, G_depth=2, dim_z=128, bottom_width=4, reso...
method init_weights (line 243) | def init_weights(self):
method forward (line 265) | def forward(self, z, y):
class DBlock (line 283) | class DBlock(nn.Module):
method __init__ (line 284) | def __init__(self, in_channels, out_channels, which_conv=layers.SNConv...
method shortcut (line 308) | def shortcut(self, x):
method forward (line 315) | def forward(self, x):
function D_arch (line 331) | def D_arch(ch=64, attention='64',ksize='333333', dilation='111111'):
class Discriminator (line 359) | class Discriminator(nn.Module):
method __init__ (line 361) | def __init__(self, D_ch=64, D_wide=True, D_depth=2, resolution=128,
method init_weights (line 458) | def init_weights(self):
method forward (line 475) | def forward(self, x, y=None):
class G_D (line 492) | class G_D(nn.Module):
method __init__ (line 493) | def __init__(self, G, D):
method forward (line 498) | def forward(self, z, gy, x=None, dy=None, train_G=False, return_G_z=Fa...
FILE: BigGAN_utils/TFHub/biggan_v1.py
function l2normalize (line 12) | def l2normalize(v, eps=1e-4):
function truncated_z_sample (line 16) | def truncated_z_sample(batch_size, z_dim, truncation=0.5, seed=None):
function denorm (line 22) | def denorm(x):
class SpectralNorm (line 27) | class SpectralNorm(nn.Module):
method __init__ (line 28) | def __init__(self, module, name='weight', power_iterations=1):
method _update_u_v (line 36) | def _update_u_v(self):
method _made_params (line 50) | def _made_params(self):
method _make_params (line 59) | def _make_params(self):
method forward (line 76) | def forward(self, *args):
class SelfAttention (line 81) | class SelfAttention(nn.Module):
method __init__ (line 84) | def __init__(self, in_dim, activation=F.relu):
method forward (line 98) | def forward(self, x):
class ConditionalBatchNorm2d (line 115) | class ConditionalBatchNorm2d(nn.Module):
method __init__ (line 116) | def __init__(self, num_features, num_classes, eps=1e-4, momentum=0.1):
method forward (line 123) | def forward(self, x, y):
class GBlock (line 131) | class GBlock(nn.Module):
method __init__ (line 132) | def __init__(
method forward (line 168) | def forward(self, input, condition=None):
class Generator128 (line 197) | class Generator128(nn.Module):
method __init__ (line 198) | def __init__(self, code_dim=120, n_class=1000, chn=96, debug=False):
method forward (line 226) | def forward(self, input, class_id):
class Generator256 (line 244) | class Generator256(nn.Module):
method __init__ (line 245) | def __init__(self, code_dim=140, n_class=1000, chn=96, debug=False):
method forward (line 272) | def forward(self, input, class_id):
class Generator512 (line 290) | class Generator512(nn.Module):
method __init__ (line 291) | def __init__(self, code_dim=128, n_class=1000, chn=96, debug=False):
method forward (line 321) | def forward(self, input, class_id):
class Discriminator (line 339) | class Discriminator(nn.Module):
method __init__ (line 340) | def __init__(self, n_class=1000, chn=96, debug=False):
method forward (line 376) | def forward(self, input, class_id):
FILE: BigGAN_utils/TFHub/converter.py
function dump_tfhub_to_hdf5 (line 43) | def dump_tfhub_to_hdf5(module_path, hdf5_path, redownload=False):
class TFHub2Pytorch (line 74) | class TFHub2Pytorch(object):
method __init__ (line 91) | def __init__(self, state_dict, tf_weights, resolution=256, load_ema=Tr...
method load (line 100) | def load(self):
method load_generator (line 104) | def load_generator(self):
method load_linear (line 119) | def load_linear(self, name_pth, name_tf, bias=True):
method load_snlinear (line 124) | def load_snlinear(self, name_pth, name_tf, bias=True):
method load_colorize (line 131) | def load_colorize(self, name_pth, name_tf):
method load_GBlock (line 134) | def load_GBlock(self, name_pth, name_tf):
method load_convs (line 138) | def load_convs(self, name_pth, name_tf):
method load_snconv (line 143) | def load_snconv(self, name_pth, name_tf, bias=True):
method load_conv (line 152) | def load_conv(self, name_pth, name_tf, bias=True):
method load_HyperBNs (line 160) | def load_HyperBNs(self, name_pth, name_tf):
method load_ScaledCrossReplicaBNs (line 164) | def load_ScaledCrossReplicaBNs(self, name_pth, name_tf):
method load_HyperBN (line 172) | def load_HyperBN(self, name_pth, name_tf):
method load_attention (line 191) | def load_attention(self, name_pth, name_tf):
method load_tf_tensor (line 199) | def load_tf_tensor(self, prefix, var, device='0'):
function convert_from_v1 (line 204) | def convert_from_v1(hub_dict, resolution=128):
function get_config (line 283) | def get_config(resolution):
function convert_biggan (line 304) | def convert_biggan(resolution, weight_dir, redownload=False, no_ema=Fals...
function generate_sample (line 328) | def generate_sample(G, z_dim, batch_size, filename, parallel=False):
function parse_args (line 342) | def parse_args():
FILE: BigGAN_utils/binary_utils.py
class BinaryActivation (line 5) | class BinaryActivation(Function):
method forward (line 8) | def forward(ctx, x):
method backward (line 13) | def backward(ctx, grad_output):
FILE: BigGAN_utils/calculate_inception_moments.py
function prepare_parser (line 19) | def prepare_parser():
function run (line 49) | def run(config):
function main (line 82) | def main():
FILE: BigGAN_utils/datasets.py
function is_image_file (line 20) | def is_image_file(filename):
function find_classes (line 33) | def find_classes(dir):
function make_dataset (line 40) | def make_dataset(dir, class_to_idx):
function pil_loader (line 58) | def pil_loader(path):
function accimage_loader (line 65) | def accimage_loader(path):
function default_loader (line 74) | def default_loader(path):
class ImageFolder (line 82) | class ImageFolder(data.Dataset):
method __init__ (line 107) | def __init__(self, root, transform=None, target_transform=None,
method __getitem__ (line 143) | def __getitem__(self, index):
method __len__ (line 166) | def __len__(self):
method __repr__ (line 169) | def __repr__(self):
class ILSVRC_HDF5 (line 184) | class ILSVRC_HDF5(data.Dataset):
method __init__ (line 185) | def __init__(self, root, transform=None, target_transform=None,
method __getitem__ (line 208) | def __getitem__(self, index):
method __len__ (line 238) | def __len__(self):
class CIFAR10 (line 243) | class CIFAR10(dset.CIFAR10):
method __init__ (line 245) | def __init__(self, root, train=True,
method __getitem__ (line 326) | def __getitem__(self, index):
method __len__ (line 347) | def __len__(self):
class CIFAR100 (line 351) | class CIFAR100(CIFAR10):
FILE: BigGAN_utils/inception_tf13.py
function prepare_parser (line 29) | def prepare_parser():
function run (line 44) | def run(config):
function main (line 130) | def main():
FILE: BigGAN_utils/inception_utils.py
class WrapInception (line 27) | class WrapInception(nn.Module):
method __init__ (line 28) | def __init__(self, net):
method forward (line 35) | def forward(self, x):
function torch_cov (line 89) | def torch_cov(m, rowvar=False):
function sqrt_newton_schulz (line 124) | def sqrt_newton_schulz(A, numIters, dtype=None):
function numpy_calculate_frechet_distance (line 144) | def numpy_calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
function torch_calculate_frechet_distance (line 200) | def torch_calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
function calculate_inception_score (line 235) | def calculate_inception_score(pred, num_splits=10):
function accumulate_inception_activations (line 248) | def accumulate_inception_activations(sample, net, num_inception_images=5...
function load_inception_net (line 261) | def load_inception_net(parallel=False):
function prepare_inception_metrics (line 274) | def prepare_inception_metrics(dataset, parallel, no_fid=False):
FILE: BigGAN_utils/layers.py
function proj (line 16) | def proj(x, y):
function gram_schmidt (line 21) | def gram_schmidt(x, ys):
function power_iteration (line 28) | def power_iteration(W, u_, update=True, eps=1e-12):
class identity (line 54) | class identity(nn.Module):
method forward (line 55) | def forward(self, input):
class SN (line 60) | class SN(object):
method __init__ (line 61) | def __init__(self, num_svs, num_itrs, num_outputs, transpose=False, ep...
method u (line 77) | def u(self):
method sv (line 83) | def sv(self):
method W_ (line 87) | def W_(self):
class SNConv2d (line 103) | class SNConv2d(nn.Conv2d, SN):
method __init__ (line 104) | def __init__(self, in_channels, out_channels, kernel_size, stride=1,
method forward (line 110) | def forward(self, x):
class SNLinear (line 116) | class SNLinear(nn.Linear, SN):
method __init__ (line 117) | def __init__(self, in_features, out_features, bias=True,
method forward (line 121) | def forward(self, x):
class SNEmbedding (line 128) | class SNEmbedding(nn.Embedding, SN):
method __init__ (line 129) | def __init__(self, num_embeddings, embedding_dim, padding_idx=None,
method forward (line 137) | def forward(self, x):
class Attention (line 144) | class Attention(nn.Module):
method __init__ (line 145) | def __init__(self, ch, which_conv=SNConv2d, name='attention'):
method forward (line 156) | def forward(self, x, y=None):
function fused_bn (line 173) | def fused_bn(x, mean, var, gain=None, bias=None, eps=1e-5):
function manual_bn (line 191) | def manual_bn(x, gain=None, bias=None, return_mean_var=False, eps=1e-5):
class myBN (line 212) | class myBN(nn.Module):
method __init__ (line 213) | def __init__(self, num_channels, eps=1e-5, momentum=0.1):
method reset_stats (line 229) | def reset_stats(self):
method forward (line 234) | def forward(self, x, gain, bias):
function groupnorm (line 259) | def groupnorm(x, norm_style):
class ccbn (line 278) | class ccbn(nn.Module):
method __init__ (line 279) | def __init__(self, output_size, input_size, which_linear, eps=1e-5, mo...
method forward (line 306) | def forward(self, x, y):
method extra_repr (line 326) | def extra_repr(self):
class bn (line 333) | class bn(nn.Module):
method __init__ (line 334) | def __init__(self, output_size, eps=1e-5, momentum=0.1,
method forward (line 359) | def forward(self, x, y=None):
class GBlock (line 375) | class GBlock(nn.Module):
method __init__ (line 376) | def __init__(self, in_channels, out_channels,
method forward (line 398) | def forward(self, x, y):
class DBlock (line 412) | class DBlock(nn.Module):
method __init__ (line 413) | def __init__(self, in_channels, out_channels, which_conv=SNConv2d, wid...
method shortcut (line 431) | def shortcut(self, x):
method forward (line 444) | def forward(self, x):
FILE: BigGAN_utils/losses.py
function loss_dcgan_dis (line 5) | def loss_dcgan_dis(dis_fake, dis_real):
function loss_dcgan_gen (line 11) | def loss_dcgan_gen(dis_fake):
function loss_hinge_dis (line 17) | def loss_hinge_dis(dis_fake, dis_real):
function loss_hinge_gen (line 27) | def loss_hinge_gen(dis_fake):
FILE: BigGAN_utils/make_hdf5.py
function prepare_parser (line 20) | def prepare_parser():
function run (line 45) | def run(config):
function main (line 102) | def main():
FILE: BigGAN_utils/sample.py
function run (line 24) | def run(config):
function main (line 174) | def main():
FILE: BigGAN_utils/sync_batchnorm/batchnorm.py
function _sum_ft (line 24) | def _sum_ft(tensor):
function _unsqueeze_ft (line 29) | def _unsqueeze_ft(tensor):
class _SynchronizedBatchNorm (line 38) | class _SynchronizedBatchNorm(_BatchNorm):
method __init__ (line 39) | def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
method forward (line 48) | def forward(self, input, gain=None, bias=None):
method __data_parallel_replicate__ (line 110) | def __data_parallel_replicate__(self, ctx, copy_id):
method _data_parallel_master (line 120) | def _data_parallel_master(self, intermediates):
method _compute_mean_std (line 147) | def _compute_mean_std(self, sum_, ssum, size):
class SynchronizedBatchNorm1d (line 162) | class SynchronizedBatchNorm1d(_SynchronizedBatchNorm):
method _check_input_dim (line 218) | def _check_input_dim(self, input):
class SynchronizedBatchNorm2d (line 225) | class SynchronizedBatchNorm2d(_SynchronizedBatchNorm):
method _check_input_dim (line 281) | def _check_input_dim(self, input):
class SynchronizedBatchNorm3d (line 288) | class SynchronizedBatchNorm3d(_SynchronizedBatchNorm):
method _check_input_dim (line 345) | def _check_input_dim(self, input):
FILE: BigGAN_utils/sync_batchnorm/batchnorm_reimpl.py
class BatchNorm2dReimpl (line 18) | class BatchNorm2dReimpl(nn.Module):
method __init__ (line 27) | def __init__(self, num_features, eps=1e-5, momentum=0.1):
method reset_running_stats (line 39) | def reset_running_stats(self):
method reset_parameters (line 43) | def reset_parameters(self):
method forward (line 48) | def forward(self, input_):
FILE: BigGAN_utils/sync_batchnorm/comm.py
class FutureResult (line 18) | class FutureResult(object):
method __init__ (line 21) | def __init__(self):
method put (line 26) | def put(self, result):
method get (line 32) | def get(self):
class SlavePipe (line 46) | class SlavePipe(_SlavePipeBase):
method run_slave (line 49) | def run_slave(self, msg):
class SyncMaster (line 56) | class SyncMaster(object):
method __init__ (line 67) | def __init__(self, master_callback):
method __getstate__ (line 78) | def __getstate__(self):
method __setstate__ (line 81) | def __setstate__(self, state):
method register_slave (line 84) | def register_slave(self, identifier):
method run_master (line 102) | def run_master(self, master_msg):
method nr_slaves (line 136) | def nr_slaves(self):
FILE: BigGAN_utils/sync_batchnorm/replicate.py
class CallbackContext (line 23) | class CallbackContext(object):
function execute_replication_callbacks (line 27) | def execute_replication_callbacks(modules):
class DataParallelWithCallback (line 50) | class DataParallelWithCallback(DataParallel):
method replicate (line 64) | def replicate(self, module, device_ids):
function patch_replication_callback (line 70) | def patch_replication_callback(data_parallel):
FILE: BigGAN_utils/sync_batchnorm/unittest.py
class TorchTestCase (line 15) | class TorchTestCase(unittest.TestCase):
method assertTensorClose (line 16) | def assertTensorClose(self, x, y):
FILE: BigGAN_utils/train.py
function run (line 34) | def run(config):
function main (line 219) | def main():
FILE: BigGAN_utils/train_fns.py
function dummy_training_function (line 14) | def dummy_training_function():
function GAN_training_function (line 20) | def GAN_training_function(G, D, GD, z_, y_, ema, state_dict, config):
function save_and_sample (line 98) | def save_and_sample(G, D, G_ema, z_, y_, fixed_z, fixed_y,
function test (line 164) | def test(G, D, G_ema, z_, y_, state_dict, config, sample, get_inception_...
FILE: BigGAN_utils/utils.py
function prepare_parser (line 30) | def prepare_parser():
function add_sample_parser (line 368) | def add_sample_parser(parser):
class CenterCropLongEdge (line 436) | class CenterCropLongEdge(object):
method __call__ (line 443) | def __call__(self, img):
method __repr__ (line 452) | def __repr__(self):
class RandomCropLongEdge (line 455) | class RandomCropLongEdge(object):
method __call__ (line 462) | def __call__(self, img):
method __repr__ (line 477) | def __repr__(self):
class MultiEpochSampler (line 483) | class MultiEpochSampler(torch.utils.data.Sampler):
method __init__ (line 492) | def __init__(self, data_source, num_epochs, start_itr=0, batch_size=128):
method __iter__ (line 503) | def __iter__(self):
method __len__ (line 521) | def __len__(self):
function get_data_loaders (line 526) | def get_data_loaders(dataset, data_root=None, augment=False, batch_size=64,
function seed_rng (line 589) | def seed_rng(seed):
function update_config_roots (line 597) | def update_config_roots(config):
function prepare_root (line 606) | def prepare_root(config):
class ema (line 616) | class ema(object):
method __init__ (line 617) | def __init__(self, source, target, decay=0.9999, start_itr=0):
method update (line 632) | def update(self, itr=None):
function ortho (line 648) | def ortho(model, strength=1e-4, blacklist=[]):
function default_ortho (line 663) | def default_ortho(model, strength=1e-4, blacklist=[]):
function toggle_grad (line 676) | def toggle_grad(model, on_or_off):
function join_strings (line 684) | def join_strings(base_string, strings):
function save_weights (line 689) | def save_weights(G, D, state_dict, weights_root, experiment_name,
function load_weights (line 714) | def load_weights(G, D, state_dict, weights_root, experiment_name,
class MetricsLogger (line 746) | class MetricsLogger(object):
method __init__ (line 747) | def __init__(self, fname, reinitialize=False):
method log (line 755) | def log(self, record=None, **kwargs):
class MyLogger (line 773) | class MyLogger(object):
method __init__ (line 774) | def __init__(self, fname, reinitialize=False, logstyle='%3.3f'):
method reinit (line 783) | def reinit(self, item):
method log (line 795) | def log(self, itr, **kwargs):
function write_metadata (line 813) | def write_metadata(logs_root, experiment_name, config, state_dict):
function progress (line 828) | def progress(items, desc='', total=None, min_delay=0.1, displaytype='s1k'):
function sample (line 867) | def sample(G, z_, y_, config):
function sample_sheet (line 879) | def sample_sheet(G, classes_per_sheet, num_classes, samples_per_class, p...
function interp (line 914) | def interp(x0, x1, num_midpoints):
function interp_sheet (line 921) | def interp_sheet(G, num_per_sheet, num_midpoints, num_classes, parallel,
function print_grad_norms (line 958) | def print_grad_norms(net):
function get_SVs (line 971) | def get_SVs(net, prefix):
function name_from_config (line 979) | def name_from_config(config):
function hashname (line 1028) | def hashname(name):
function query_gpu (line 1039) | def query_gpu(indices):
function count_parameters (line 1044) | def count_parameters(module):
function sample_1hot (line 1050) | def sample_1hot(batch_size, num_classes, device='cuda'):
class Distribution (line 1062) | class Distribution(torch.Tensor):
method init_distribution (line 1064) | def init_distribution(self, dist_type, **kwargs):
method sample_ (line 1072) | def sample_(self):
method to (line 1081) | def to(self, *args, **kwargs):
function prepare_z_y (line 1089) | def prepare_z_y(G_batch_size, dim_z, nclasses, device='cuda',
function initiate_standing_stats (line 1104) | def initiate_standing_stats(net):
function accumulate_standing_stats (line 1111) | def accumulate_standing_stats(net, z, y, nclasses, num_accumulations=16):
class Adam16 (line 1131) | class Adam16(Optimizer):
method __init__ (line 1132) | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,weigh...
method load_state_dict (line 1139) | def load_state_dict(self, state_dict):
method step (line 1147) | def step(self, closure=None):
FILE: DiffAugment_pytorch.py
function DiffAugment (line 10) | def DiffAugment(x, policy='', channels_first=True):
function rand_brightness (line 23) | def rand_brightness(x):
function rand_saturation (line 28) | def rand_saturation(x):
function rand_contrast (line 34) | def rand_contrast(x):
function rand_translation (line 40) | def rand_translation(x, ratio=0.125): ### ratio: org: 0.125
function rand_resize (line 55) | def rand_resize(x, min_ratio=0.8, max_ratio=1.2): ### ratio: org: 0.125
function rand_cutout (line 80) | def rand_cutout(x, ratio=0.5):
FILE: fusedream_utils.py
function AugmentLoss (line 19) | def AugmentLoss(img, clip_model, text, replicate=10, interp_mode='biline...
function NaiveSemanticLoss (line 33) | def NaiveSemanticLoss(img, clip_model, text, interp_mode='bilinear'):
function get_gaussian_mask (line 46) | def get_gaussian_mask(size=256):
function save_image (line 59) | def save_image(img, path, n_per_row=1):
function get_G (line 68) | def get_G(resolution=256):
class FuseDreamBaseGenerator (line 144) | class FuseDreamBaseGenerator():
method __init__ (line 145) | def __init__(self, G, G_config, G_batch_size=10, clip_mode="ViT-B/32",...
method generate_basis (line 170) | def generate_basis(self, text, init_iters=500, num_basis=5):
method optimize_clip_score (line 218) | def optimize_clip_score(self, z_init_cllt, y_init_cllt, text, latent_n...
method measureAugCLIP (line 291) | def measureAugCLIP(self, z, y, text, augment=False, num_samples=20):
Condensed preview — 48 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (322K chars).
[
{
"path": "BigGAN_utils/BigGAN.py",
"chars": 20897,
"preview": "import numpy as np\nimport math\nimport functools\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport tor"
},
{
"path": "BigGAN_utils/BigGANdeep.py",
"chars": 22982,
"preview": "import numpy as np\nimport math\nimport functools\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport tor"
},
{
"path": "BigGAN_utils/LICENSE",
"chars": 1067,
"preview": "MIT License\n\nCopyright (c) 2019 Andy Brock\n\nPermission is hereby granted, free of charge, to any person obtaining a copy"
},
{
"path": "BigGAN_utils/README.md",
"chars": 150,
"preview": "Download pre-trained weights from (https://drive.google.com/drive/folders/1nJ3HmgYgeA9NZr-oU-enqbYeO7zBaANs?usp=sharing)"
},
{
"path": "BigGAN_utils/TFHub/biggan_v1.py",
"chars": 12174,
"preview": "# BigGAN V1:\n# This is now deprecated code used for porting the TFHub modules to pytorch,\n# included here for reference "
},
{
"path": "BigGAN_utils/TFHub/converter.py",
"chars": 17423,
"preview": "\"\"\"Utilities for converting TFHub BigGAN generator weights to PyTorch.\nRecommended usage:\nTo convert all BigGAN variants"
},
{
"path": "BigGAN_utils/__init__.py",
"chars": 46,
"preview": "import sys\nsys.path.append('./BigGAN_utils/')\n"
},
{
"path": "BigGAN_utils/animal_hash.py",
"chars": 32285,
"preview": "c = ['Aardvark', 'Abyssinian', 'Affenpinscher', 'Akbash', 'Akita', 'Albatross',\n 'Alligator', 'Alpaca', 'Angelfish',"
},
{
"path": "BigGAN_utils/binary_utils.py",
"chars": 303,
"preview": "from torch.autograd import Function\nfrom torch.optim import SGD\n\n\nclass BinaryActivation(Function):\n\n @staticmethod\n "
},
{
"path": "BigGAN_utils/calculate_inception_moments.py",
"chars": 3551,
"preview": "''' Calculate Inception Moments\n This script iterates over the dataset and calculates the moments of the \n activations o"
},
{
"path": "BigGAN_utils/datasets.py",
"chars": 11778,
"preview": "''' Datasets\r\n This file contains definitions for our CIFAR, ImageFolder, and HDF5 datasets\r\n'''\r\nimport os\r\nimport o"
},
{
"path": "BigGAN_utils/inception_tf13.py",
"chars": 5363,
"preview": "''' Tensorflow inception score code\nDerived from https://github.com/openai/improved-gan\nCode derived from tensorflow/ten"
},
{
"path": "BigGAN_utils/inception_utils.py",
"chars": 12310,
"preview": "''' Inception utilities\n This file contains methods for calculating IS and FID, using either\n the original numpy c"
},
{
"path": "BigGAN_utils/layers.py",
"chars": 17130,
"preview": "''' Layers\n This file contains various layers for the BigGAN models.\n'''\nimport numpy as np\nimport torch\nimport torch"
},
{
"path": "BigGAN_utils/logs/BigGAN_ch96_bs256x8.jsonl",
"chars": 8920,
"preview": "{\"itr\": 2000, \"IS_mean\": 2.806771755218506, \"IS_std\": 0.019480662420392036, \"FID\": 173.76484159711126, \"_stamp\": 1551403"
},
{
"path": "BigGAN_utils/logs/compare_IS.m",
"chars": 2196,
"preview": "clc\nclear all\nclose all\nfclose all;\n\n\n\n%% Get All logs and sort them\ns = {};\nd = dir();\nj = 1;\nfor i = 1:length(d)\n i"
},
{
"path": "BigGAN_utils/logs/metalog.txt",
"chars": 3888,
"preview": "datetime: 2019-03-18 13:27:59.181225\nconfig: {'dataset': 'I128_hdf5', 'augment': False, 'num_workers': 8, 'pin_memory': "
},
{
"path": "BigGAN_utils/logs/process_inception_log.m",
"chars": 445,
"preview": "function [itr, IS, FID, t] = process_inception_log(fname, which_log)\nf = sprintf('%s_%s',fname, which_log);%'G_loss.log'"
},
{
"path": "BigGAN_utils/logs/process_training.m",
"chars": 3293,
"preview": "clc\nclear all\nclose all\nfclose all;\n\n\n\n%% Get all training logs for a given run\ntarget_dir = '.';\ns = {};\nnm = {};\nd = d"
},
{
"path": "BigGAN_utils/losses.py",
"chars": 821,
"preview": "import torch\nimport torch.nn.functional as F\n\n# DCGAN loss\ndef loss_dcgan_dis(dis_fake, dis_real):\n L1 = torch.mean(F.s"
},
{
"path": "BigGAN_utils/make_hdf5.py",
"chars": 4971,
"preview": "\"\"\" Convert dataset to HDF5\n This script preprocesses a dataset and saves it (images and labels) to \n an HDF5 file"
},
{
"path": "BigGAN_utils/sample.py",
"chars": 8346,
"preview": "''' Sample\n This script loads a pretrained net and a weightsfile and sample '''\nimport functools\nimport math\nimport nu"
},
{
"path": "BigGAN_utils/scripts/launch_BigGAN_bs256x8.sh",
"chars": 640,
"preview": "#!/bin/bash\npython train.py \\\n--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 256 --load_in_mem "
},
{
"path": "BigGAN_utils/scripts/launch_BigGAN_bs512x4.sh",
"chars": 640,
"preview": "#!/bin/bash\npython train.py \\\n--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 512 --load_in_mem "
},
{
"path": "BigGAN_utils/scripts/launch_BigGAN_ch64_bs256x8.sh",
"chars": 638,
"preview": "#!/bin/bash\npython train.py \\\n--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 256 --load_in_mem "
},
{
"path": "BigGAN_utils/scripts/launch_BigGAN_deep.sh",
"chars": 671,
"preview": "#!/bin/bash\npython train.py \\\n--model BigGANdeep \\\n--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_siz"
},
{
"path": "BigGAN_utils/scripts/launch_SAGAN_bs128x2_ema.sh",
"chars": 535,
"preview": "#!/bin/bash\npython train.py \\\n--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 128 \\\n--num_G_accu"
},
{
"path": "BigGAN_utils/scripts/launch_SNGAN.sh",
"chars": 507,
"preview": "#!/bin/bash\npython train.py \\\n--dataset I128_hdf5 --parallel --shuffle --num_workers 8 --batch_size 64 \\\n--num_G_accum"
},
{
"path": "BigGAN_utils/scripts/launch_cifar_ema.sh",
"chars": 406,
"preview": "#!/bin/bash\nCUDA_VISIBLE_DEVICES=0,1 python train.py \\\n--shuffle --batch_size 50 --parallel \\\n--num_G_accumulations 1 --"
},
{
"path": "BigGAN_utils/scripts/sample_BigGAN_bs256x8.sh",
"chars": 936,
"preview": "# use z_var to change the variance of z for all the sampling\n# use --mybn --accumulate_stats --num_standing_accumulation"
},
{
"path": "BigGAN_utils/scripts/sample_cifar_ema.sh",
"chars": 426,
"preview": "#!/bin/bash\nCUDA_VISIBLE_DEVICES=0,1 python sample.py \\\n--shuffle --batch_size 50 --G_batch_size 256 --parallel \\\n--num_"
},
{
"path": "BigGAN_utils/scripts/utils/duplicate.sh",
"chars": 941,
"preview": "#duplicate.sh\nsource=BigGAN_I128_hdf5_seed0_Gch64_Dch64_bs256_Glr1.0e-04_Dlr4.0e-04_Gnlinplace_relu_Dnlinplace_relu_Gini"
},
{
"path": "BigGAN_utils/scripts/utils/prepare_data.sh",
"chars": 155,
"preview": "#!/bin/bash\npython make_hdf5.py --dataset I128 --batch_size 256 --data_root data\npython calculate_inception_moments.py -"
},
{
"path": "BigGAN_utils/sync_batchnorm/__init__.py",
"chars": 449,
"preview": "# -*- coding: utf-8 -*-\n# File : __init__.py\n# Author : Jiayuan Mao\n# Email : maojiayuan@gmail.com\n# Date : 27/01/2"
},
{
"path": "BigGAN_utils/sync_batchnorm/batchnorm.py",
"chars": 14882,
"preview": "# -*- coding: utf-8 -*-\n# File : batchnorm.py\n# Author : Jiayuan Mao\n# Email : maojiayuan@gmail.com\n# Date : 27/01/"
},
{
"path": "BigGAN_utils/sync_batchnorm/batchnorm_reimpl.py",
"chars": 2383,
"preview": "#! /usr/bin/env python3\n# -*- coding: utf-8 -*-\n# File : batchnorm_reimpl.py\n# Author : acgtyrant\n# Date : 11/01/201"
},
{
"path": "BigGAN_utils/sync_batchnorm/comm.py",
"chars": 4449,
"preview": "# -*- coding: utf-8 -*-\n# File : comm.py\n# Author : Jiayuan Mao\n# Email : maojiayuan@gmail.com\n# Date : 27/01/2018\n"
},
{
"path": "BigGAN_utils/sync_batchnorm/replicate.py",
"chars": 3226,
"preview": "# -*- coding: utf-8 -*-\n# File : replicate.py\n# Author : Jiayuan Mao\n# Email : maojiayuan@gmail.com\n# Date : 27/01/"
},
{
"path": "BigGAN_utils/sync_batchnorm/unittest.py",
"chars": 746,
"preview": "# -*- coding: utf-8 -*-\n# File : unittest.py\n# Author : Jiayuan Mao\n# Email : maojiayuan@gmail.com\n# Date : 27/01/2"
},
{
"path": "BigGAN_utils/train.py",
"chars": 9153,
"preview": "\"\"\" BigGAN: The Authorized Unofficial PyTorch release\n Code by A. Brock and A. Andonian\n This code is an unofficia"
},
{
"path": "BigGAN_utils/train_fns.py",
"chars": 8584,
"preview": "''' train_fns.py\nFunctions for the main loop of training different conditional image models\n'''\nimport torch\nimport torc"
},
{
"path": "BigGAN_utils/utils.py",
"chars": 48779,
"preview": "#!/usr/bin/env python\n# -*- coding: utf-8 -*-\n\n''' Utilities file\nThis file contains utility functions for bookkeeping, "
},
{
"path": "BigGAN_utils/weights/README.md",
"chars": 120,
"preview": "Download pre-trained weights from \nhttps://drive.google.com/drive/folders/1nJ3HmgYgeA9NZr-oU-enqbYeO7zBaANs?usp=sharing\n"
},
{
"path": "DiffAugment_pytorch.py",
"chars": 4134,
"preview": "# Differentiable Augmentation for Data-Efficient GAN Training\n# Shengyu Zhao, Zhijian Liu, Ji Lin, Jun-Yan Zhu, and Song"
},
{
"path": "LICENSE",
"chars": 1065,
"preview": "MIT License\n\nCopyright (c) 2021 gnobitab\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\no"
},
{
"path": "README.md",
"chars": 2278,
"preview": "# FuseDream\n\nThis repo contains code for our paper ([paper link](https://arxiv.org/abs/2112.01573)):\n\n**FuseDream: Train"
},
{
"path": "fusedream_generator.py",
"chars": 1264,
"preview": "import torch\nfrom tqdm import tqdm\nfrom torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize\ni"
},
{
"path": "fusedream_utils.py",
"chars": 11892,
"preview": "import torch\nfrom tqdm import tqdm\nfrom torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize\ni"
}
]
About this extraction
This page contains the full source code of the gnobitab/FuseDream GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 48 files (303.0 KB), approximately 88.8k tokens, and a symbol index with 296 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.