Kaggle competition: Generative Dog Images
<SOS> metric: PostProcessorKernel Kaggle data source: generative-dog-images<define_search_model>
import datetime
import gc
import json
import math
import os
import time
import shutil
import sys
from functools import partial
from numba import jit, njit
from PIL import Image
from tqdm import tqdm, trange
import cv2
import numpy as np
import matplotlib.pyplot as plt
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 torch.utils.data as data
from torch.utils.data import DataLoader
import torchvision
from torchvision import datasets
import torchvision.transforms as T
import albumentations as A
from albumentations.pytorch import ToTensor
import xml.etree.ElementTree as ET
class Attention(nn.Module):
def __init__(self, channels, reduction_attn=8, reduction_sc=2):
super().__init__()
self.channles_attn = channels // reduction_attn
self.channels_sc = channels // reduction_sc
self.conv_query = spectral_norm(nn.Conv2d(channels, self.channles_attn, kernel_size=1, bias=False))
self.conv_key = spectral_norm(nn.Conv2d(channels, self.channles_attn, kernel_size=1, bias=False))
self.conv_value = spectral_norm(nn.Conv2d(channels, self.channels_sc, kernel_size=1, bias=False))
self.conv_attn = spectral_norm(nn.Conv2d(self.channels_sc, channels, kernel_size=1, bias=False))
self.gamma = nn.Parameter(torch.zeros(1))
nn.init.orthogonal_(self.conv_query.weight.data)
 nn.init.orthogonal_(self.conv_key.weight.data)
 nn.init.orthogonal_(self.conv_value.weight.data)
 nn.init.orthogonal_(self.conv_attn.weight.data)
 def forward(self, x):
batch, _, h, w = x.size()
proj_query = self.conv_query(x ).view(batch, self.channles_attn, -1)
 proj_key = F.max_pool2d(self.conv_key(x), 2 ).view(batch, self.channles_attn, -1)
 attn = torch.bmm(proj_key.permute(0,2,1), proj_query)
 attn = F.softmax(attn, dim=1)
 proj_value = F.max_pool2d(self.conv_value(x), 2 ).view(batch, self.channels_sc, -1)
 attn = torch.bmm(proj_value, attn)
 attn = attn.view(batch, self.channels_sc, h, w)
 attn = self.conv_attn(attn)
 out = self.gamma * attn + x
return out<train_on_grid>
INPUT_PATH = '.. /input/all-dogs'
ANNOT_PATH = '.. /input/annotation/Annotation'
OUT_DIR = '.. /output_images'
IMG_SIZE = 64
IMG_SIZE_2 = IMG_SIZE * 2
IMG_CHANNELS = 3
IMG_EXTENSIONS =('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp' )
class CBN2d(nn.Module):
def __init__(self, num_features, num_conditions):
super().__init__()
self.bn = nn.BatchNorm2d(num_features, affine=False)
 self.embed = spectral_norm(nn.Conv2d(num_conditions, num_features*2, kernel_size=1, bias=False))
nn.init.orthogonal_(self.embed.weight.data)
 def forward(self, x, y):
out = self.bn(x)
 embed = self.embed(y.unsqueeze(2 ).unsqueeze(3))
gamma, beta = embed.chunk(2, dim=1)
 out =(1.0 + gamma)* out + beta
return out<define_search_model>
def proj(x, y):
return torch.mm(y, x.t())* y / torch.mm(y, y.t())
 def gram_schmidt(x, ys):
for y in ys:
x = x - proj(x, y)
 return x
def power_iteration(W, u_, update=True, eps=1e-12):
us, vs, svs = [], [], []
for i, u in enumerate(u_):
with torch.no_grad() :
v = torch.matmul(u, W)
 v = F.normalize(gram_schmidt(v, vs), eps=eps)
 vs += [v]
u = torch.matmul(v, W.t())
 u = F.normalize(gram_schmidt(u, us), eps=eps)
 us += [u]
if update:
u_[i][:] = u
svs += [torch.squeeze(torch.matmul(torch.matmul(v, W.t()), u.t())) ]
return svs, us, vs
class identity(nn.Module):
def forward(self, input):
return input
class SN(object):
def __init__(self, num_svs, num_itrs, num_outputs, transpose=False, eps=1e-12):
self.num_itrs = num_itrs
self.num_svs = num_svs
self.transpose = transpose
self.eps = eps
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))
@property
def u(self):
return [getattr(self, 'u%d' % i)for i in range(self.num_svs)]
@property
def sv(self):
return [getattr(self, 'sv%d' % i)for i in range(self.num_svs)]
def W_(self):
W_mat = self.weight.view(self.weight.size(0), -1)
 if self.transpose:
W_mat = W_mat.t()
for _ in range(self.num_itrs):
svs, us, vs = power_iteration(W_mat, self.u, update=self.training, eps=self.eps)
 if self.training:
with torch.no_grad() :
for i, sv in enumerate(svs):
self.sv[i][:] = sv
return self.weight / svs[0]
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)
 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)
 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_())
 class Attention(nn.Module):
def __init__(self, ch, which_conv=SNConv2d, name='attention'):
super(Attention, self ).__init__()
self.ch = ch
self.which_conv = which_conv
self.theta = self.which_conv(self.ch, self.ch // 4, kernel_size=1, padding=0, bias=False)
 self.phi = self.which_conv(self.ch, self.ch // 4, 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)
 self.gamma = P(torch.tensor(0.) , requires_grad=True)
 def forward(self, x, y=None):
theta = self.theta(x)
 phi = F.max_pool2d(self.phi(x), [2, 2])
 g = F.max_pool2d(self.g(x), [2, 2])
 theta = theta.view(-1, self.ch // 4, x.shape[2] * x.shape[3])
 phi = phi.view(-1, self.ch // 4, x.shape[2] * x.shape[3] // 4)
 g = g.view(-1, self.ch // 2, x.shape[2] * x.shape[3] // 4)
 beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), -1)
 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
def fused_bn(x, mean, var, gain=None, bias=None, eps=1e-5):
scale = torch.rsqrt(var + eps)
 if gain is not None:
scale = scale * gain
shift = mean * scale
if bias is not None:
shift = shift - bias
return x * scale - shift
def manual_bn(x, gain=None, bias=None, return_mean_var=False, eps=1e-5):
float_x = x.float()
m = torch.mean(float_x, [0, 2, 3], keepdim=True)
 m2 = torch.mean(float_x ** 2, [0, 2, 3], keepdim=True)
 var =(m2 - m ** 2)
 var = var.type(x.type())
 m = m.type(x.type())
 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)
 class myBN(nn.Module):
def __init__(self, num_channels, eps=1e-5, momentum=0.1):
super(myBN, self ).__init__()
self.momentum = momentum
self.eps = eps
self.momentum = momentum
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))
self.accumulate_standing = False
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 self.accumulate_standing:
self.stored_mean[:] = self.stored_mean + mean.data
self.stored_var[:] = self.stored_var + var.data
self.accumulation_counter += 1.0
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
else:
mean = self.stored_mean.view(1, -1, 1, 1)
 var = self.stored_var.view(1, -1, 1, 1)
 if self.accumulate_standing:
mean = mean / self.accumulation_counter
var = var / self.accumulation_counter
return fused_bn(x, mean, var, gain, bias, self.eps)
 def groupnorm(x, norm_style):
if 'ch' in norm_style:
ch = int(norm_style.split('_')[-1])
 groups = max(int(x.shape[1])// ch, 1)
 elif 'grp' in norm_style:
groups = int(norm_style.split('_')[-1])
 else:
groups = 16
return F.group_norm(x, groups)
 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
self.gain = which_linear(input_size, output_size)
 self.bias = which_linear(input_size, output_size)
 self.eps = eps
self.momentum = momentum
self.cross_replica = cross_replica
self.mybn = mybn
self.norm_style = norm_style
if 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):
gain =(1 + self.gain(y)).view(y.size(0), -1, 1, 1)
 bias = self.bias(y ).view(y.size(0), -1, 1, 1)
 if self.mybn:
return self.bn(x, gain=gain, bias=bias)
 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__)
 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
self.gain = P(torch.ones(output_size), requires_grad=True)
 self.bias = P(torch.zeros(output_size), requires_grad=True)
 self.eps = eps
self.momentum = momentum
self.cross_replica = cross_replica
self.mybn = mybn
if mybn:
self.bn = myBN(output_size, self.eps, self.momentum)
 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.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)
 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
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)
 self.bn1 = self.which_bn(in_channels)
 self.bn2 = self.which_bn(out_channels)
 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
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
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
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 = 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 )
class GBlock(nn.Module):
def __init__(self, in_channels, out_channels, num_conditions, upsample=False):
super().__init__()
self.upsample = upsample
self.learnable_sc = in_channels != out_channels or upsample
self.conv1 = spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False))
self.conv2 = spectral_norm(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False))
self.cbn1 = CBN2d(in_channels, num_conditions)
 self.cbn2 = CBN2d(out_channels, num_conditions)
 if self.learnable_sc:
self.conv_sc = spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False))
self.relu = nn.ReLU()
nn.init.orthogonal_(self.conv1.weight.data)
 nn.init.orthogonal_(self.conv2.weight.data)
 if self.learnable_sc:
nn.init.orthogonal_(self.conv_sc.weight.data)
 def _upsample_conv(self, x, conv):
x = F.interpolate(x, scale_factor=2, mode='nearest')
 x = conv(x)
 return x
def _residual(self, x, y):
x = self.relu(self.cbn1(x, y))
x = self._upsample_conv(x, self.conv1)if self.upsample else self.conv1(x)
 x = self.relu(self.cbn2(x, y))
x = self.conv2(x)
 return x
def _shortcut(self, x):
if self.learnable_sc:
x = self._upsample_conv(x, self.conv_sc)if self.upsample else self.conv_sc(x)
 return x
def forward(self, x, y):
return self._shortcut(x)+ self._residual(x, y )<choose_model_class>
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__()
self.ch = G_ch
self.dim_z = dim_z
self.bottom_width = bottom_width
self.resolution = resolution
self.kernel_size = G_kernel_size
self.attention = G_attn
self.n_classes = n_classes
self.G_shared = G_shared
self.shared_dim = shared_dim if shared_dim > 0 else dim_z
self.hier = hier
self.cross_replica = cross_replica
self.mybn = mybn
self.activation = G_activation
self.init = G_init
self.G_param = G_param
self.norm_style = norm_style
self.BN_eps = BN_eps
self.SN_eps = SN_eps
self.fp16 = G_fp16
self.arch = G_arch(self.ch, self.attention)[resolution]
if self.hier:
self.num_slots = len(self.arch['in_channels'])+ 1
self.z_chunk_size =(self.dim_z // self.num_slots)
 self.dim_z = self.z_chunk_size * self.num_slots
else:
self.num_slots = 1
self.z_chunk_size = 0
if self.G_param == 'SN':
self.which_conv = partial(SNConv2d,
kernel_size=3, padding=1,
num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
eps=self.SN_eps)
 self.which_linear = partial(SNLinear,
num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,
eps=self.SN_eps)
 else:
self.which_conv = partial(nn.Conv2d, kernel_size=3, padding=1)
 self.which_linear = nn.Linear
self.which_embedding = nn.Embedding
bn_linear =(partial(self.which_linear, bias=False)if self.G_shared else self.which_embedding)
 self.which_bn = partial(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)
 self.shared =(self.which_embedding(n_classes, self.shared_dim)if G_shared else identity())
 self.linear = self.which_linear(
 self.dim_z // self.num_slots,
self.arch['in_channels'][0] *(self.bottom_width ** 2))
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['out_channels'][index],
which_conv=self.which_conv,
which_bn=self.which_bn,
activation=self.activation,
upsample=(partial(F.interpolate, scale_factor=2)if self.arch['upsample'][index] else None)
)]]
if self.arch['attention'][self.arch['resolution'][index]]:
print('Adding attention layer in G at resolution %d' % self.arch['resolution'][index])
 self.blocks[-1] += [Attention(self.arch['out_channels'][index], self.which_conv)]
self.blocks = nn.ModuleList([nn.ModuleList(block)for block in self.blocks])
 self.output_layer = nn.Sequential(
 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)
)
 if not skip_init:
self.init_weights()
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...')
 self.optim = 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)
 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)
 def forward(self, z, y):
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)
 h = self.linear(z)
 h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width)
 for index, blocklist in enumerate(self.blocks):
for block in blocklist:
h = block(h, ys[index])
 return torch.tanh(self.output_layer(h))
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__()
self.ch = D_ch
self.D_wide = D_wide
self.resolution = resolution
self.kernel_size = D_kernel_size
self.attention = D_attn
self.n_classes = n_classes
self.activation = D_activation
self.init = D_init
self.D_param = D_param
self.SN_eps = SN_eps
self.fp16 = D_fp16
self.arch = D_arch(self.ch, self.attention)[resolution]
if self.D_param == 'SN':
self.which_conv = partial(SNConv2d,
kernel_size=3, padding=1,
num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
 self.which_linear = partial(SNLinear,
num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
 self.which_embedding = partial(SNEmbedding,
num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,
eps=self.SN_eps)
 self.blocks = []
for index in range(len(self.arch['out_channels'])) :
self.blocks += [[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 self.arch['attention'][self.arch['resolution'][index]]:
print('Adding attention layer in D at resolution %d' % self.arch['resolution'][index])
 self.blocks[-1] += [Attention(self.arch['out_channels'][index], self.which_conv)]
self.blocks = nn.ModuleList([nn.ModuleList(block)for block in self.blocks])
 self.linear = self.which_linear(self.arch['out_channels'][-1], output_dim)
 self.embed = self.which_embedding(self.n_classes, self.arch['out_channels'][-1])
 if not skip_init:
self.init_weights()
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...')
 self.optim = 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)
 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):
h = x
for index, blocklist in enumerate(self.blocks):
for block in blocklist:
h = block(h)
 h = torch.sum(self.activation(h), [2, 3])
 out = self.linear(h)
 out = out + torch.sum(self.embed(y)* h, 1, keepdim=True)
 return out
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):
with torch.set_grad_enabled(train_G):
G_z = self.G(z, self.G.shared(gy))
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()
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
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
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]])
 else:
if return_G_z:
return D_out, G_z
else:
return D_out
class Generator(nn.Module):
def __init__(self, latent_dim, ch, num_classes, embed_dim, use_attn=False):
super().__init__()
self.latent_dim = latent_dim
self.ch = ch
self.num_classes = num_classes
self.embed_dim = embed_dim
self.use_attn = use_attn
self.num_chunk = 5
num_latents = self.__get_num_latents()
self.embed = nn.Embedding(num_classes, embed_dim)
 self.fc = spectral_norm(nn.Linear(num_latents[0], ch*8*4*4, bias=False))
self.block1 = GBlock(ch*8, ch*8, num_latents[1], upsample=True)
 self.block2 = GBlock(ch*8, ch*4, num_latents[2], upsample=True)
 self.block3 = GBlock(ch*4, ch*2, num_latents[3], upsample=True)
 if use_attn:
self.attn = Attention(ch*2)
 self.block4 = GBlock(ch*2, ch, num_latents[4], upsample=True)
 self.bn = nn.BatchNorm2d(ch)
 self.relu = nn.ReLU()
self.conv_last = spectral_norm(nn.Conv2d(ch, 3, kernel_size=3, padding=1, bias=False))
self.tanh = nn.Tanh()
nn.init.orthogonal_(self.embed.weight.data)
 nn.init.orthogonal_(self.fc.weight.data)
 nn.init.orthogonal_(self.conv_last.weight.data)
 nn.init.constant_(self.bn.weight.data, 1.0)
 nn.init.constant_(self.bn.bias.data, 0.0)
 def __get_num_latents(self):
xs = torch.empty(self.latent_dim ).chunk(self.num_chunk)
 num_latents = [x.size(0)for x in xs]
for i in range(1, self.num_chunk):
num_latents[i] += self.embed_dim
return num_latents
def forward(self, x, y):
xs = x.chunk(self.num_chunk, dim=1)
 y = self.embed(y)
 h = self.fc(xs[0])
 h = h.view(h.size(0), self.ch*8, 4, 4)
 h = self.block1(h, torch.cat([y, xs[1]], dim=1))
h = self.block2(h, torch.cat([y, xs[2]], dim=1))
h = self.block3(h, torch.cat([y, xs[3]], dim=1))
if self.use_attn:
h = self.attn(h)
 h = self.block4(h, torch.cat([y, xs[4]], dim=1))
h = self.relu(self.bn(h))
out = self.tanh(self.conv_last(h))
return out<define_search_model>
class MyImg() :
def __init__(self, img, tfm):
self.px = np.array(img)
 self.tfm = tfm
@property
def size(self):
h, w, _ = self.px.shape
return min(w, h)
 def pad(img, padding_mode='reflect'):
p = math.ceil(( max(img.size)- min(img.size)) / 2)
 p_horr = p if img.width < img.height else 0
p_vert = p if img.height < img.width else 0
img = T.Pad(( p_horr, p_vert), padding_mode=padding_mode )(img)
 if img.width != img.height:
s = min(img.size)
 img = img.crop(( 0, 0, s, s))
return img
def take_top(img):
size = min(img.size)
 bbox =(0, 0, size, size)
 return img.crop(bbox)
 def take_diagonal(img):
w, h = img.size
size = min(w, h)
 bbox_l =(0, 0, size, size)
 bbox_r =(w - size, h - size, w, h)
 return [img.crop(bbox_l), img.crop(bbox_r)]
resize = T.Resize(IMG_SIZE, interpolation=Image.LANCZOS)
 resize2x = T.Resize(IMG_SIZE_2, interpolation=Image.LANCZOS)
 top_crop = T.Compose([T.Lambda(take_top), resize])
 top_crop2x = T.Compose([T.Lambda(take_top), resize2x])
 two_crops = T.Compose([resize, T.Lambda(take_diagonal)])
 two_crops2x = T.Compose([resize2x, T.Lambda(take_diagonal)])
 pad_only = T.Compose([T.Lambda(pad), resize])
 pad_only2x = T.Compose([T.Lambda(pad), resize2x])
 @jit(nopython=True)
 def pad_one_dim(clow, chigh, pad, cmax):
clow = max(0, clow - pad)
 chigh = min(cmax, chigh + pad)
 return clow, chigh, chigh - clow
def calc_bbox(obj, img_w, img_h, zoom=0.0, try_square=True):
bndbox = obj.find('bndbox')
 xmin = int(bndbox.find('xmin' ).text)
 ymin = int(bndbox.find('ymin' ).text)
 xmax = int(bndbox.find('xmax' ).text)
 ymax = int(bndbox.find('ymax' ).text)
 xmin, xmax, obj_w = pad_one_dim(xmin, xmax, 0, img_w)
 ymin, ymax, obj_h = pad_one_dim(ymin, ymax, 0, img_h)
 if zoom != 0.0:
pad_w = obj_w * zoom / 2
pad_h = obj_h * zoom / 2
xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad_w, img_w)
 ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad_h, img_h)
 if try_square:
if obj_w > obj_h:
pad =(obj_w - obj_h)/ 2
ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad, img_h)
 elif obj_h > obj_w:
pad =(obj_h - obj_w)/ 2
xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad, img_w)
 if obj_w > obj_h:
pad = obj_w - obj_h
ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad, img_h)
 elif obj_h > obj_w:
pad = obj_h - obj_w
xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad, img_w)
 return int(xmin), int(ymin), int(xmax), int(ymax)
 @jit(nopython=True)
 def bb2wh(bbox):
width = bbox[2] - bbox[0]
height = bbox[3] - bbox[1]
return width, height
def create_counters() :
return { 'top': 0, 'sides': 0, 'slide': 0, 'pad': 0, 'skip': 0, 'noop': 0 }
def make_x1res(img, bbox, counter):
ar = img.width / img.height
if ar == 1.0:
tfm_imgs = resize(img); counter['noop'] += 1
elif 1.0 < ar < 1.4 or 1.0 < 1/ar < 1.25:
tfm_imgs = pad_only(img); counter['pad'] += 1
elif 1.4 <= ar < 2.0:
tfm_imgs = two_crops(img); counter['sides'] += 2
elif 1.25 <= 1/ar < 1.5:
tfm_imgs = [pad_only(img), top_crop(img)]
counter['pad'] += 1; counter['top'] += 1
elif 1.5 <= 1/ar < 1.8:
tfm_imgs = top_crop(img); counter['top'] += 1
else:
tfm_imgs = None; counter['skip'] += 1
return tfm_imgs
def make_x2res(img, bbox, counter):
ar = img.width / img.height
if ar == 1.0:
tfm_imgs = resize2x(img); counter['noop'] += 1
elif 1.0 < ar < 1.4 or 1.0 < 1/ar < 1.4:
tfm_imgs = pad_only2x(img); counter['pad'] += 1
else:
tfm_imgs = None
return tfm_imgs
def add_sample(samples, label, tfm, imgs, labels):
if not samples: return
elif isinstance(samples, Image.Image):
imgs.append(MyImg(samples, tfm))
labels.append(label)
 elif isinstance(samples, list):
imgs.extend([MyImg(s, tfm)for s in samples])
 labels.extend([label] * len(samples))
else: assert False
def is_valid_file(x):
return datasets.folder.has_file_allowed_extension(x, IMG_EXTENSIONS)
 class DogsDataSet(datasets.vision.VisionDataset):
def __init__(self, root, transforms, target_transform=None, max_samples=None):
super().__init__(root, transform=None)
 assert isinstance(transforms, list)and len(transforms)== 4
self.transforms = transforms
self.target_transform = target_transform
self.max_samples = max_samples
self.classes = {}
imgs, labels = self._load_subfolders_images(self.root)
 assert len(imgs)== len(labels)
 if len(imgs)== 0:
raise RuntimeError(f'Found 0 files in subfolders of: {self.root}')
 self.imgs = imgs
self.labels = labels
def _create_or_get_class(self, name):
try:
label = self.classes[name]
except KeyError:
label = len(self.classes)
 self.classes[name] = label
return label
def _load_subfolders_images(self, root):
imgs, labels, paths = [], [], []
add_sample_ = partial(add_sample, imgs=imgs, labels=labels)
 low_res = lambda bbox: min(bb2wh(bbox)) < IMG_SIZE
can_2x = lambda bbox: min(bb2wh(bbox)) >= IMG_SIZE_2
counter = create_counters()
for root, _, fnames in sorted(os.walk(root)) :
for fname in sorted(fnames):
path = os.path.join(root, fname)
 paths.append(path)
 if self.max_samples:
paths = paths[:self.max_samples]
for path in paths:
if not is_valid_file(path): continue
img = datasets.folder.default_loader(path)
 annotation_basename = os.path.splitext(os.path.basename(path)) [0]
annotation_dirname = next(dirname for dirname in os.listdir(ANNOT_PATH)if dirname.startswith(annotation_basename.split('_')[0]))
annotation_filename = os.path.join(ANNOT_PATH, annotation_dirname, annotation_basename)
 tree = ET.parse(annotation_filename)
 root = tree.getroot()
objects = root.findall('object')
 for o in objects:
name = o.find('name' ).text
label = self._create_or_get_class(name)
 prev_bbox, tfm_imgs = None, None
bbox = calc_bbox(o, img_w=img.width, img_h=img.height, zoom=0.03125)
 if low_res(bbox):
counter['skip'] += 1; continue
obj_img = img.crop(bbox)
 ar = obj_img.width / obj_img.height
tfm_imgs = make_x1res(obj_img, bbox, counter)
 add_sample_(tfm_imgs, label, 0)
 if can_2x(bbox):
tfm_imgs = make_x2res(obj_img, bbox, counter)
 add_sample_(tfm_imgs, label, 1)
 prev_bbox = bbox
bbox = calc_bbox(o, img_w=img.width, img_h=img.height, zoom=0.1)
 if not(bbox == prev_bbox or low_res(bbox)) :
obj_img = img.crop(bbox)
 ar = obj_img.width / obj_img.height
if 1.05 < ar < 1.25:
tfm_imgs = resize(obj_img); counter['slide'] += 1
add_sample_(tfm_imgs, label, 3)
 if can_2x(bbox):
tfm_imgs = make_x2res(obj_img, bbox, counter)
 add_sample_(tfm_imgs, label, 2)
 bbox = calc_bbox(o, img_w=img.width, img_h=img.height, zoom=0.12, try_square=False)
 if not low_res(bbox):
obj_img = img.crop(bbox)
 ar = obj_img.width / obj_img.height
if 1.3 <= 1/ar < 1.8:
tfm_imgs = top_crop(obj_img); counter['top'] += 1
add_sample_(tfm_imgs, label, 0)
 if can_2x(bbox):
tfm_imgs = make_x2res(obj_img, bbox, counter)
 add_sample_(tfm_imgs, label, 1)
 n_x1, n_x2 = 0, 0
for i, img in enumerate(imgs):
if img.size == IMG_SIZE:
n_x1 +=1
else:
n_x2 +=1
print(f'Found {len(self.classes)} classes
Loaded 64x64 {n_x1} images
'
f'Loaded 128x128 {n_x2} images
')
 print(counter)
 return imgs, labels
def __getitem__(self, index):
img = self.imgs[index]
label = self.labels[index]
tfms = self.transforms[img.tfm]
img = tfms(image=img.px)['image']
if self.target_transform:
label = self.target_transform(label)
 return img, label
def __len__(self):
return len(self.imgs)
 def create_runtime_tfms() :
mean, std = [0.5]*3, [0.5]*3
resize_to_64 = A.SmallestMaxSize(IMG_SIZE, interpolation=cv2.INTER_LANCZOS4)
 out = [A.HorizontalFlip(p=0.5), A.Normalize(mean=mean, std=std), ToTensor() ]
rand_crop = A.RandomCrop(IMG_SIZE, IMG_SIZE)
 affine_1 = A.ShiftScaleRotate(
 shift_limit=0, scale_limit=0.1, rotate_limit=8,
interpolation=cv2.INTER_CUBIC,
border_mode=cv2.BORDER_REFLECT_101, p=1.0)
 affine_1 = A.Compose([affine_1, resize_to_64])
 affine_2 = A.ShiftScaleRotate(
 shift_limit=0.06, scale_limit=(-0.06, 0.14), rotate_limit=6,
interpolation=cv2.INTER_CUBIC,
border_mode=cv2.BORDER_REFLECT_101, p=1.0)
 affine_2 = A.Compose([affine_2, resize_to_64])
 tfm_0 = A.Compose(out)
 tfm_1 = A.Compose([affine_1, *out])
 tfm_2 = A.Compose([affine_2, *out])
 tfm_3 = A.Compose([rand_crop, *out])
 return [tfm_0, tfm_1, tfm_2, tfm_3]
def get_data_loaders(data_root=None, batch_size=32, num_workers=2, shuffle=True,
pin_memory=True, drop_last=True):
print('Using dataset root location %s' % data_root)
 train_set = DogsDataSet(data_root, create_runtime_tfms())
 loaders = []
loader_kwargs = {'num_workers': num_workers, 'pin_memory': pin_memory,
'drop_last': drop_last}
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=shuffle, **loader_kwargs)
 loaders.append(train_loader)
 return loaders
class DBlock(nn.Module):
def __init__(self, in_channels, out_channels, downsample=False, optimized=False):
super().__init__()
self.downsample = downsample
self.optimized = optimized
self.learnable_sc = in_channels != out_channels or downsample
self.conv1 = spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False))
self.conv2 = spectral_norm(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False))
if self.learnable_sc:
self.conv_sc = spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False))
self.relu = nn.ReLU()
nn.init.orthogonal_(self.conv1.weight.data)
 nn.init.orthogonal_(self.conv2.weight.data)
 if self.learnable_sc:
nn.init.orthogonal_(self.conv_sc.weight.data)
 def _residual(self, x):
if not self.optimized:
x = self.relu(x)
 x = self.conv1(x)
 x = self.relu(x)
 x = self.conv2(x)
 if self.downsample:
x = F.avg_pool2d(x, 2)
 return x
def _shortcut(self, x):
if self.learnable_sc:
if self.optimized:
x = self.conv_sc(F.avg_pool2d(x, 2)) if self.downsample else self.conv_sc(x)
 else:
x = F.avg_pool2d(self.conv_sc(x), 2)if self.downsample else self.conv_sc(x)
 return x
def forward(self, x):
return self._shortcut(x)+ self._residual(x )<choose_model_class>
def seed_rng(seed):
torch.manual_seed(seed)
 torch.cuda.manual_seed(seed)
 np.random.seed(seed)
 def update_config_roots(config):
if config['base_root']:
print('Pegging all root folders to base root %s' % config['base_root'])
 for key in ['data', 'weights', 'logs', 'samples']:
config['%s_root' % key] = '%s/%s' %(config['base_root'], key)
 return config
def prepare_root(config):
for key in ['weights_root', 'logs_root', 'samples_root']:
if not os.path.exists(config[key]):
print('Making directory %s for %s...' %(config[key], key))
os.mkdir(config[key])
 def join_strings(base_string, strings):
return base_string.join([item for item in strings if item] )
class Discriminator(nn.Module):
def __init__(self, ch, num_classes, use_attn=False):
super().__init__()
self.ch = ch
self.num_classes = num_classes
self.use_attn = use_attn
self.block1 = DBlock(3, ch, downsample=True, optimized=True)
 if use_attn:
self.attn = Attention(ch)
 self.block2 = DBlock(ch, ch*2, downsample=True)
 self.block3 = DBlock(ch*2, ch*4, downsample=True)
 self.block4 = DBlock(ch*4, ch*8, downsample=True)
 self.block5 = DBlock(ch*8+1, ch*8, downsample=False)
 self.relu = nn.ReLU()
self.fc = spectral_norm(nn.Linear(ch*8, 1, bias=False))
self.embed = spectral_norm(nn.Embedding(num_classes, ch*8))
self.clf = spectral_norm(nn.Linear(ch*8, num_classes, bias=False))
nn.init.orthogonal_(self.fc.weight.data)
 nn.init.orthogonal_(self.embed.weight.data)
 nn.init.orthogonal_(self.clf.weight.data)
 def minibatch_stddev(self, x, group_size=4, eps=1e-8):
shape = x.size()
y = x.view(group_size, -1, shape[1], shape[2], shape[3])
 y -= torch.mean(y, dim=0, keepdim=True)
 y = torch.mean(y.pow(2), dim=0)
 y = torch.sqrt(y + eps)
 y = torch.mean(y, dim=[1,2,3], keepdim=True)
 y = y.repeat(group_size, 1, shape[2], shape[3])
 return torch.cat([x, y], dim=1)
 def forward(self, x, y):
h = self.block1(x)
 if self.use_attn:
h = self.attn(h)
 h = self.block2(h)
 h = self.block3(h)
 h = self.block4(h)
 h = self.block5(self.minibatch_stddev(h))
h = self.relu(h)
 h = torch.sum(h, dim=(2,3))
out = self.fc(h)
 out += torch.sum(self.embed(y)*h, dim=1, keepdim=True)
 ac = self.clf(h)
 ac = F.log_softmax(ac, dim=1)
 return out, ac<categorify>
def run(G, D, GD, loaders, state_dict, config, device, G_ema=None, ema=None):
G_batch_size = max(config['G_batch_size'], config['batch_size'])
 z_, y_ = prepare_z_y(
 G_batch_size, G.dim_z, config['n_classes'], device=device, fp16=config['G_fp16'])
 fixed_z, fixed_y = prepare_z_y(
 G_batch_size, G.dim_z, config['n_classes'], device=device, fp16=config['G_fp16'])
 fixed_z.sample_()
fixed_y.sample_()
train = create_train_fn(G, D, GD, z_, y_, ema, state_dict, config)
 print('Beginning training at epoch %d...' % state_dict['epoch'])
 start_time = time.perf_counter()
total_iters = config['num_epochs'] * len(loaders[0])
 for epoch in range(state_dict['epoch'], config['num_epochs']):
for i,(x, y)in enumerate(loaders[0]):
state_dict['itr'] += 1
G.train()
D.train()
if config['ema']:
G_ema.train()
if config['D_fp16']:
x, y = x.to(device ).half() , y.to(device)
 else:
x, y = x.to(device), y.to(device)
 metrics = train(x, y)
 if not(state_dict['itr'] % config['log_interval']):
curr_time = time.perf_counter()
curr_time_str = datetime.datetime.fromtimestamp(curr_time ).strftime('%H:%M:%S')
 elapsed = str(datetime.timedelta(seconds=(curr_time - start_time)))
 log =(
 "[{}] [{}] [{} / {}] Ep {}, ".format(curr_time_str, elapsed, state_dict['itr'], total_iters, epoch)+
', '.join(['%s : %+4.3f' %(key, metrics[key])for key in metrics])
)
 print(log)
 if not(state_dict['itr'] % config['save_every']):
if config['G_eval_mode']:
print('Switchin G to eval mode...')
 G.eval()
save_and_sample(G, G_ema, fixed_z, fixed_y, config)
 if config['stop_after'] > 0 and int(time.perf_counter() - start_time)> config['stop_after']:
print("Time limit reached! Stopping training!")
 return
state_dict['epoch'] += 1
netG = Generator(**config['Generator'] ).to(device, torch.float32)
 netD = Discriminator(**config['Discriminator'] ).to(device, torch.float32 )<load_pretrained>
def trunc_trick(bs, z_dim, bound=2.0):
z = torch.randn(bs, z_dim)
 while z.min() < -bound or bound < z.max() :
z = z.where(( -bound < z)&(z < bound), torch.randn_like(z))
return z
def collect_bn_stats(G, n_samples, config):
im_batch_size = config['n_classes']
G.train()
for i_batch in range(0, n_samples, im_batch_size):
with torch.no_grad() :
z = torch.randn(im_batch_size, G.dim_z, device=device)
 y = torch.arange(im_batch_size ).to(device)
 images = G(z, G.shared(y)).float().cpu()
def generate_images(out_dir, G, n_images, config):
im_batch_size = config['n_classes']
z_bound = config['trunc_z']
if z_bound > 0.0: print(f'Truncating z to(-{z_bound}, {z_bound})')
 for i_batch in range(0, n_images, im_batch_size):
with torch.no_grad() :
if z_bound > 0.0:
z = trunc_trick(im_batch_size, G.dim_z, bound=z_bound ).to(device)
 else:
z = torch.randn(im_batch_size, G.dim_z, device=device)
 y = torch.arange(im_batch_size ).to(device)
 images = G(z, G.shared(y)).float().cpu()
if i_batch + im_batch_size > n_images:
n_last_images = n_images - i_batch
print(f'Taking only {n_last_images} images from the last batch...')
 images = images[:n_last_images]
for i_image, image in enumerate(images):
fname = os.path.join(out_dir, f'image_{i_batch+i_image:05d}.png')
 image = denorm(image)
 if config['denoise']:
image = image * 256
image = image.numpy().transpose(1, 2, 0 ).astype('uint8')
 image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
 image = cv2.fastNlMeansDenoisingColored(
 src=image, dst=None,
h=config['denoise_str_lum'],
hColor=config['denoise_str_col'],
templateWindowSize=config['denoise_kernel_size'],
searchWindowSize=config['denoise_search_window'] 
)
 cv2.imwrite(fname, image)
 else:
torchvision.utils.save_image(image, fname)
 def show_saved_samples(n_images=256, cols=8, size=3):
rows = math.ceil(n_images / cols)
 _, axes = plt.subplots(rows, cols, figsize=(cols * size, rows * size))
for i, ax in enumerate(axes.flatten()):
img = Image.open(f'{OUT_DIR}/image_{i:05d}.png')
 ax.get_xaxis().set_visible(False)
 ax.get_yaxis().set_visible(False)
 ax.imshow(img)
 plt.subplots_adjust(wspace=0.0, hspace=0.0)
 plt.show()
netGE = Generator(**config['Generator'] ).to(device, torch.float32)
 netGE.load_state_dict(netG.state_dict());<choose_model_class>
activation_dict = {
'inplace_relu': nn.ReLU(inplace=True),
'relu': nn.ReLU(inplace=False),
'ir': nn.ReLU(inplace=True),
}
config = {
'num_workers': 2,
'pin_memory': True,
'shuffle': True,
'use_multiepoch_sampler': False,
'model': 'BigGAN',
'G_param': 'SN',
'D_param': 'SN',
'G_ch': 64,
'D_ch': 64,
'G_depth': 1,
'D_depth': 1,
'D_wide': True,
'G_shared': False,
'shared_dim': 0,
'dim_z': 128,
'z_var': 1.0,
'hier': False,
'cross_replica': False,
'mybn': False,
'G_nl': 'relu',
'D_nl': 'relu',
'G_attn': '64',
'D_attn': '64',
'norm_style': 'bn',
'seed': 0,
'G_init': 'ortho',
'D_init': 'ortho',
'skip_init': False,
'G_lr': 5e-5,
'D_lr': 2e-4,
'G_B1': 0.0,
'D_B1': 0.0,
'G_B2': 0.999,
'D_B2': 0.999,
'batch_size': 64,
'G_batch_size': 0,
'num_G_accumulations': 1,
'num_D_steps': 2,
'num_D_accumulations': 1,
'split_D': False,
'num_epochs': 100,
'parallel': False,
'G_fp16': False,
'D_fp16': False,
'D_mixed_precision': False,
'G_mixed_precision': False,
'accumulate_stats': False,
'num_standing_accumulations': 16,
'G_eval_mode': False,
'save_every': 2000,
'num_save_copies': 2,
'num_best_copies': 2,
'which_best': 'IS',
'no_fid': False,
'test_every': 5000,
'num_inception_images': 50000,
'hashname': False,
'base_root': '',
'data_root': '.. /data',
'weights_root': '.. /weights',
'logs_root': '.. /logs',
'samples_root': '.. /samples',
'pbar': 'mine',
'name_suffix': '',
'experiment_name': '',
'config_from_name': False,
'ema': False,
'ema_decay': 0.9999,
'use_ema': False,
'ema_start': 0,
'adam_eps': 1e-8,
'BN_eps': 1e-5,
'SN_eps': 1e-8,
'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,
'log_interval': 100,
}
config['resolution'] = 64
config['n_classes'] = 120
config['data_root'] = INPUT_PATH
optim_G = Adam(params=netG.parameters() , lr=config['lr_G'], betas=config['betas'])
 optim_D = Adam(params=netD.parameters() , lr=config['lr_D'], betas=config['betas'] )<choose_model_class>
config['num_epochs'] = 100
config['num_D_steps'] = 1
config['G_lr'] = 1e-4 * 0.85
config['D_lr'] = 4e-4 * 0.85
config['D_B2'] = 0.999
config['G_B2'] = 0.999
config['G_ch'] = 32
config['D_ch'] = 64
config['G_attn'] = '32'
config['D_attn'] = '32'
config['G_nl'] = 'inplace_relu'
config['D_nl'] = 'inplace_relu'
config['G_shared'] = True
config['shared_dim'] = 128
config['dim_z'] = 120
config['hier'] = True
config['SN_eps'] = 1e-8
config['BN_eps'] = 1e-5
config['adam_eps'] = 1e-8
config['G_ortho'] = 0.0
config['G_init'] = 'ortho'
config['D_init'] = 'ortho'
config['G_eval_mode'] = True
config['save_every'] = 5000
config['log_interval'] = 500
config['batch_size'] = 32
config['ema'] = True
config['use_ema'] = True
config['ema_decay'] = 0.9999
config['ema_start'] = 50_000
config['stop_after'] = 31_500
config['trunc_z'] = 0.0
config['denoise'] = True
config['denoise_str_lum'] = 2.65
config['denoise_str_col'] = 1.0
config['denoise_kernel_size'] = 7
config['denoise_search_window'] = 11
decay_iter = config['num_iterations'] - config['decay_start_iteration']
if decay_iter > 0:
lr_lambda_G = lambda x:(max(0,1-x/decay_iter))
lr_lambda_D = lambda x:(max(0,1-x/(decay_iter*config['d_steps'])))
 lr_sche_G = LambdaLR(optim_G, lr_lambda=lr_lambda_G)
 lr_sche_D = LambdaLR(optim_D, lr_lambda=lr_lambda_D )<compute_test_metric>
config = update_config_roots(config)
 device = 'cuda'
seed_rng(1234)
 prepare_root(config)
 torch.backends.cudnn.benchmark = True
def calc_advloss_D(real, fake, margin=1.0):
loss_real = torch.mean(( real - fake.mean() - margin)** 2)
 loss_fake = torch.mean(( fake - real.mean() + margin)** 2)
 loss =(loss_real + loss_fake)/ 2
return loss<compute_test_metric>
%%time
D_batch_size =(config['batch_size'] * config['num_D_steps'] * config['num_D_accumulations'])
 loaders = get_data_loaders(
 data_root=INPUT_PATH,
batch_size=D_batch_size,
num_workers=config['num_workers'],
shuffle=config['shuffle'],
pin_memory=config['pin_memory'],
drop_last=True
)
def calc_advloss_G(real, fake, margin=1.0):
loss_real = torch.mean(( real - fake.mean() + margin)** 2)
 loss_fake = torch.mean(( fake - real.mean() - margin)** 2)
 loss =(loss_real + loss_fake)/ 2
return loss<choose_model_class>
config['G_activation'] = activation_dict[config['G_nl']]
config['D_activation'] = activation_dict[config['D_nl']]
G = Generator(**config ).to(device)
 D = Discriminator(**config ).to(device)
 if config['ema']:
print('Preparing EMA for G with decay of {}'.format(config['ema_decay']))
G_ema = Generator(**{**config, 'skip_init': True, 'no_optim': True} ).to(device)
 ema = apply_ema(G, G_ema, config['ema_decay'], config['ema_start'])
 else:
G_ema, ema = None, None
GD = G_D(G, D)
 print('Number of params in G: {} D: {}'.format(
 *[sum([p.data.nelement() for p in net.parameters() ])for net in [G, D]]))
criterion = nn.NLLLoss().to(device, torch.float32 )<define_variables>
state_dict = {
'itr': 0,
'epoch': 0,
'config': config
}
def sample_latents(batch_size, latent_dim, num_classes):
latents = torch.randn(( batch_size, latent_dim), dtype=torch.float32, device=device)
 labels = torch.randint(0, num_classes, size=(batch_size,), dtype=torch.long, device=device)
 return latents, labels<categorify>
run(G, D, GD, loaders, state_dict, config, device, G_ema=G_ema, ema=ema )
step = 1
while True:
for i in range(config['d_steps']):
for param in netD.parameters() :
param.requires_grad_(True)
 optim_D.zero_grad()
real_imgs, real_labels = train_dataiterator.__next__()
batch_size = real_imgs.size(0)
 latents, fake_labels = sample_latents(batch_size, **config['sample_latents'])
 fake_imgs = netG(latents, fake_labels ).detach()
preds_real, preds_real_labels = netD(real_imgs, real_labels)
 preds_fake, _ = netD(fake_imgs, fake_labels)
 loss_D = calc_advloss_D(preds_real, preds_fake, config['margin'])
 loss_D += config['gamma'] * criterion(preds_real_labels, real_labels)
 loss_D.backward()
optim_D.step()
if(decay_iter > 0)and(step > config['decay_start_iteration']):
lr_sche_D.step()
for param in netD.parameters() :
param.requires_grad_(False)
 optim_G.zero_grad()
real_imgs, real_labels = train_dataiterator.__next__()
batch_size = real_imgs.size(0)
 latents, fake_labels = sample_latents(batch_size, **config['sample_latents'])
 fake_imgs = netG(latents, fake_labels)
 preds_real, _ = netD(real_imgs, real_labels)
 preds_fake, preds_fake_labels = netD(fake_imgs, fake_labels)
 loss_G = calc_advloss_G(preds_real, preds_fake, config['margin'])
 loss_G += config['gamma'] * criterion(preds_fake_labels, fake_labels)
 loss_G.backward()
optim_G.step()
if(decay_iter > 0)and(step > config['decay_start_iteration']):
lr_sche_G.step()
for param_G, param_GE in zip(netG.parameters() , netGE.parameters()):
param_GE.data.mul_(config['ema'] ).add_(( 1-config['ema'])*param_G.data)
 for buffer_G, buffer_GE in zip(netG.buffers() , netGE.buffers()):
buffer_GE.data.mul_(config['ema'] ).add_(( 1-config['ema'])*buffer_G.data)
 if step < config['num_iterations']:
step += 1
else:
print('total step: {}'.format(step))
break<normalization>
%%time
if not os.path.exists(OUT_DIR):
os.mkdir(OUT_DIR)
 which_G = G_ema if config['ema'] and config['use_ema'] else G
if config['use_ema']:
collect_bn_stats(G_ema, 10_000, config)
 gc.collect()
G.eval()
generate_images(OUT_DIR, which_G, 10_000, config)
 shutil.make_archive('images', 'zip', OUT_DIR )