Kaggle competition: Generative Dog Images
<SOS> metric: PostProcessorKernel Kaggle data source: generative-dog-images<choose_model_class>
from functools import partial
from multiprocessing import Pool
import os
from pathlib import Path
import random
import shutil
import time
import xml.etree.ElementTree as ET
from matplotlib import pyplot as plt
import numpy as np
from PIL import Image
from scipy.stats import truncnorm
import torch
from torch import nn, optim
from torch.optim.optimizer import Optimizer
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms
from torchvision.datasets.folder import default_loader
from torchvision.utils import save_image
from torch.nn.utils import spectral_norm
from tqdm import tqdm
from torch.autograd import Variable
import sys
def define_generator(latent_dim):
	model = Sequential()
	
	n_nodes = 128 * 7 * 7
	model.add(Dense(n_nodes, input_dim=latent_dim))
	model.add(LeakyReLU(alpha=0.2))
	model.add(Reshape(( 7, 7, 128)))
 	
	model.add(Conv2DTranspose(128,(4,4), strides=(2,2), padding='same'))
	model.add(LeakyReLU(alpha=0.2))
	
	model.add(Conv2DTranspose(128,(4,4), strides=(2,2), padding='same'))
	model.add(LeakyReLU(alpha=0.2))
	
	model.add(Conv2D(1,(7,7), activation='tanh', padding='same'))
	return model<categorify>
start_time = time.time()
batch_size = 32
epochs = 210
seed = 1029
TRAIN_DIR = Path('.. /input/all-dogs/')
 ANNOTATION_DIR = Path('.. /input/annotation/Annotation/')
 DOG_DIR = Path('.. /dogs/dogs/')
 OUT_DIR = Path('.. /output_images/')
 DOG_DIR.mkdir(parents=True, exist_ok=True)
 OUT_DIR.mkdir(exist_ok=True)
 device = torch.device('cuda')
 lr = 0.0005
beta1 = 0.5
nz = 256
real_label = 0.95
fake_label = 0
random.seed(seed)
 os.environ['PYTHONHASHSEED'] = str(seed)
 np.random.seed(seed)
 torch.manual_seed(seed)
 torch.cuda.manual_seed(seed)
 torch.backends.cudnn.deterministic = True
def define_generator(latent_dim, n_classes=120):
	
	in_label = Input(shape=(1,))
	
	li = Embedding(n_classes, 50 )(in_label)
 	
	n_nodes = 7 * 7
	li = Dense(n_nodes )(li)
 	
	li = Reshape(( 7, 7, 1))(li)
 	
	in_lat = Input(shape=(latent_dim,))
	
	n_nodes = 128 * 7 * 7
	gen = Dense(n_nodes )(in_lat)
 	gen = LeakyReLU(alpha=0.2 )(gen)
 	gen = Reshape(( 7, 7, 128))(gen)
 	
	merge = Concatenate()([gen, li])
 	
	gen = Conv2DTranspose(128,(4,4), strides=(2,2), padding='same' )(merge)
 	gen = LeakyReLU(alpha=0.2 )(gen)
 	
	gen = Conv2DTranspose(128,(4,4), strides=(2,2), padding='same' )(gen)
 	gen = LeakyReLU(alpha=0.2 )(gen)
 	
	out_layer = Conv2D(1,(7,7), activation='tanh', padding='same' )(gen)
 	
	model = Model([in_lat, in_label], out_layer)
 	return model<choose_model_class>
mylist = os.listdir(ANNOTATION_DIR )
discriminator.trainable=False
gan_input = Input(shape=(10000,))
x = generator(gan_input)
 gan_output = discriminator(x)
 gan = Model(gan_input, gan_output)
 gan.get_layer('model_1' ).get_layer('conv' ).set_weights([np.array([[[[-1 ]]],[[[255.]]]])])
 gan.compile(optimizer=Adam(5), loss='mean_squared_error')
<train_model>
label_map = {}
for listname in mylist:
label_map[listname.split('-')[0]]=listname.split('-')[1]
train = np.zeros(( 10000,10000))
for i in range(10000): train[i,i] = 1
zeros = np.zeros(( 10000,12288))
lr = 5.
 for k in range(50):
annealer = LearningRateScheduler(lambda x: lr)
 h = gan.fit(train, zeros, epochs = 1, batch_size=256, callbacks=[annealer], verbose=0)
 if(k<10)|(k%5==4):
print('Epoch',(k+1)*10,'/500 - loss =',h.history['loss'][-1])
 if h.history['loss'][-1] < 25: lr = 1.
 if h.history['loss'][-1] < 1.5: lr = 0.5
if k<10:
plt.figure(figsize=(15,3))
for j in range(5):
xx = np.zeros(( 10000))
xx[np.random.randint(10000)] = 1
plt.subplot(1,5,j+1)
 img = generator.predict(xx.reshape(( -1,10000)))[0].reshape(( -1,64,64,3))
img = Image.fromarray(( img ).astype('uint8' ).reshape(( 64,64,3)))
 plt.axis('off')
 plt.imshow(img)
 plt.show()<predict_on_test>
label_number={}
for index, label in enumerate(label_map):
label_number[label]=index
class DogGenerator:
index = 0
def getDog(self,seed):
xx = np.zeros(( 10000))
xx[self.index] = 1.0
img = generator.predict(xx.reshape(( -1,10000)))[0].reshape(( 64,64,3))
self.index =(self.index+1)%10000
return Image.fromarray(img.astype('uint8'))<choose_model_class>
class PixelwiseNorm(nn.Module):
def __init__(self):
super(PixelwiseNorm, self ).__init__()
def forward(self, x, alpha=1e-8):

 y = x.pow(2.).mean(dim=1, keepdim=True ).add(alpha ).sqrt()
y = x / y
return y
class KernelEvalException(Exception):
pass
model_params = {
'Inception': {
'name': 'Inception',
'imsize': 64,
'output_layer': 'Pretrained_Net/pool_3:0',
'input_layer': 'Pretrained_Net/ExpandDims:0',
'output_shape': 2048,
'cosine_distance_eps': 0.1
}
}
def create_model_graph(pth):

 with tf.gfile.FastGFile(pth, 'rb')as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
 _ = tf.import_graph_def(graph_def, name='Pretrained_Net')
 def _get_model_layer(sess, model_name):
layername = model_params[model_name]['output_layer']
layer = sess.graph.get_tensor_by_name(layername)
 ops = layer.graph.get_operations()
for op_idx, op in enumerate(ops):
for o in op.outputs:
shape = o.get_shape()
if shape._dims != []:
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.__dict__['_shape_val'] = tf.TensorShape(new_shape)
 return layer
def get_activations(images, sess, model_name, batch_size=50, verbose=False):

 inception_layer = _get_model_layer(sess, model_name)
 n_images = images.shape[0]
if batch_size > n_images:
print("warning: batch size is bigger than the data size.setting batch size to data size")
 batch_size = n_images
n_batches = n_images//batch_size + 1
pred_arr = np.empty(( n_images,model_params[model_name]['output_shape']))
for i in tqdm(range(n_batches)) :
if verbose:
print("\rPropagating batch %d/%d" %(i+1, n_batches), end="", flush=True)
 start = i*batch_size
if start+batch_size < n_images:
end = start+batch_size
else:
end = n_images
batch = images[start:end]
pred = sess.run(inception_layer, {model_params[model_name]['input_layer']: batch})
 pred_arr[start:end] = pred.reshape(-1,model_params[model_name]['output_shape'])
 if verbose:
print(" done")
 return pred_arr
def normalize_rows(x: np.ndarray):

 return np.nan_to_num(x/np.linalg.norm(x, ord=2, axis=1, keepdims=True))
def cosine_distance(features1, features2):
features1_nozero = features1[np.sum(features1, axis=1)!= 0]
features2_nozero = features2[np.sum(features2, axis=1)!= 0]
norm_f1 = normalize_rows(features1_nozero)
 norm_f2 = normalize_rows(features2_nozero)
 d = 1.0-np.abs(np.matmul(norm_f1, norm_f2.T))
print('d.shape=',d.shape)
 print('np.min(d, axis=1 ).shape=',np.min(d, axis=1 ).shape)
 mean_min_d = np.mean(np.min(d, axis=1))
print('distance=',mean_min_d)
 return mean_min_d
def distance_thresholding(d, eps):
if d < eps:
return d
else:
return 1
def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):

 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
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
warnings.warn(msg)
 offset = np.eye(sigma1.shape[0])* eps
covmean = linalg.sqrtm(( sigma1 + offset ).dot(sigma2 + offset))
if np.iscomplexobj(covmean):
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
print('covmean.shape=',covmean.shape)
 tr_covmean = np.trace(covmean)
 return diff.dot(diff)+ np.trace(sigma1)+ np.trace(sigma2)- 2 * tr_covmean
def calculate_activation_statistics(images, sess, model_name, batch_size=50, verbose=False):

 act = get_activations(images, sess, model_name, batch_size, verbose)
 mu = np.mean(act, axis=0)
 sigma = np.cov(act, rowvar=False)
 return mu, sigma, act
def _handle_path_memorization(path, sess, model_name, is_checksize, is_check_png):
path = pathlib.Path(path)
 files = list(path.glob('*.jpg')) + list(path.glob('*.png'))
imsize = model_params[model_name]['imsize']
x = np.array([np.array(img_read_checks(fn, imsize, is_checksize, imsize, is_check_png)) for fn in files])
 m, s, features = calculate_activation_statistics(x, sess, model_name)
 del x
return m, s, features
def img_read_checks(filename, resize_to, is_checksize=False, check_imsize = 64, is_check_png = False):
im = Image.open(str(filename))
if is_checksize and im.size !=(check_imsize,check_imsize):
raise KernelEvalException('The images are not of size '+str(check_imsize))
if is_check_png and im.format != 'PNG':
raise KernelEvalException('Only PNG images should be submitted.')
 if resize_to is None:
return im
else:
return im.resize(( resize_to,resize_to),Image.ANTIALIAS)
 def calculate_kid_given_paths(paths, model_name, model_path, feature_path=None, mm=[], ss=[], ff=[]):

 tf.reset_default_graph()
create_model_graph(str(model_path))
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
 m1, s1, features1 = _handle_path_memorization(paths[0], sess, model_name, is_checksize = True, is_check_png = True)
 if len(mm)!= 0:
m2 = mm
s2 = ss
features2 = ff
elif feature_path is None:
m2, s2, features2 = _handle_path_memorization(paths[1], sess, model_name, is_checksize = False, is_check_png = False)
 else:
with np.load(feature_path)as f:
m2, s2, features2 = f['m'], f['s'], f['features']
print('m1,m2 shape=',(m1.shape,m2.shape),'s1,s2=',(s1.shape,s2.shape))
print('starting calculating FID')
 fid_value = calculate_frechet_distance(m1, s1, m2, s2)
 print('done with FID, starting distance calculation')
 distance = cosine_distance(features1, features2)
 return fid_value, distance, m2, s2, features2<load_pretrained>
class ConditionalBatchNorm2d(nn.BatchNorm2d):

 def __init__(self, num_features, eps=1e-05, momentum=0.1,
affine=False, track_running_stats=True):
super(ConditionalBatchNorm2d, self ).__init__(
 num_features, eps, momentum, affine, track_running_stats 
)
 def forward(self, input, weight, bias, **kwargs):
self._check_input_dim(input)
 exponential_average_factor = 0.0
if self.training and self.track_running_stats:
self.num_batches_tracked += 1
if self.momentum is None:
exponential_average_factor = 1.0 / self.num_batches_tracked.item()
else:
exponential_average_factor = self.momentum
output = F.batch_norm(input, self.running_mean, self.running_var,
self.weight, self.bias,
self.training or not self.track_running_stats,
exponential_average_factor, self.eps)
 if weight.dim() == 1:
weight = weight.unsqueeze(0)
 if bias.dim() == 1:
bias = bias.unsqueeze(0)
 size = output.size()
weight = weight.unsqueeze(-1 ).unsqueeze(-1 ).expand(size)
 bias = bias.unsqueeze(-1 ).unsqueeze(-1 ).expand(size)
 return weight * output + bias
class CategoricalConditionalBatchNorm2d(ConditionalBatchNorm2d):
def __init__(self, num_classes, num_features, eps=1e-5, momentum=0.1,
affine=False, track_running_stats=True):
super(CategoricalConditionalBatchNorm2d, self ).__init__(
 num_features, eps, momentum, affine, track_running_stats 
)
 self.weights = nn.Embedding(num_classes, num_features)
 self.biases = nn.Embedding(num_classes, num_features)
 self._initialize()
def _initialize(self):
torch.nn.init.ones_(self.weights.weight.data)
 torch.nn.init.zeros_(self.biases.weight.data)
 def forward(self, input, c, **kwargs):
weight = self.weights(c)
 bias = self.biases(c)
 return super(CategoricalConditionalBatchNorm2d, self ).forward(input, weight, bias )
if ComputeLB:
with zipfile.ZipFile(".. /working/images.zip","r")as z:
z.extractall(".. /tmp/images2/")
 m2 = []; s2 =[]; f2 = []
user_images_unzipped_path = '.. /tmp/images2/'
images_path = [user_images_unzipped_path,'.. /input/generative-dog-images/all-dogs/all-dogs/']
public_path = '.. /input/dog-face-generation-competition-kid-metric-input/classify_image_graph_def.pb'
fid_epsilon = 10e-15
fid_value_public, distance_public, m2, s2, f2 = calculate_kid_given_paths(images_path, 'Inception', public_path, mm=m2, ss=s2, ff=f2)
 distance_public = distance_thresholding(distance_public, model_params['Inception']['cosine_distance_eps'])
 print("FID_public: ", fid_value_public, "distance_public: ", distance_public, "multiplied_public: ",
fid_value_public /(distance_public + fid_epsilon))
! rm -r.. /tmp<import_modules>
class ResBlockGenerator(nn.Module):
def __init__(self, in_channels, out_channels, stride=1):
super(ResBlockGenerator, self ).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, padding=1)
 self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, padding=1)
 self.model = nn.Sequential(
 nn.BatchNorm2d(in_channels),
nn.ReLU() ,
nn.Upsample(scale_factor=2),
self.conv1,
nn.BatchNorm2d(out_channels),
nn.ReLU() ,
self.conv2 
)
 self.bypass = nn.Sequential()
if stride != 1:
if in_channels != out_channels:
self.bypass = nn.Sequential(
 nn.Upsample(scale_factor=2),
nn.Conv2d(in_channels,out_channels, 1, 1, padding=0)
)
 else:
self.bypass = nn.Upsample(scale_factor=2)
 def forward(self, x):
return self.model(x)+ self.bypass(x)
 class ResBlockDiscriminator(nn.Module):
def __init__(self, in_channels, out_channels, stride=1):
super(ResBlockDiscriminator, self ).__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, 3, 1, padding=1)
 self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, padding=1)
 nn.init.xavier_uniform(self.conv1.weight.data, 1.)
nn.init.xavier_uniform(self.conv2.weight.data, 1.)
if stride == 1:
self.model = nn.Sequential(
 nn.ReLU() ,
spectral_norm(self.conv1),
nn.ReLU() ,
spectral_norm(self.conv2)
)
 else:
self.model = nn.Sequential(
 nn.ReLU() ,
spectral_norm(self.conv1),
nn.ReLU() ,
spectral_norm(self.conv2),
nn.AvgPool2d(2, stride=stride, padding=0)
)
 self.bypass = nn.Sequential()
if stride != 1:
self.bypass_conv = nn.Conv2d(in_channels,out_channels, 1, 1, padding=0)
 nn.init.xavier_uniform(self.bypass_conv.weight.data, np.sqrt(2))
self.bypass = nn.Sequential(
 spectral_norm(self.bypass_conv),
nn.AvgPool2d(2, stride=stride, padding=0)
)
 def forward(self, x):
return self.model(x)+ self.bypass(x )
batch_size = 32
start = time()
def seed_everything(seed=42):
random.seed(seed)
 os.environ['PYTHONHASHSEED'] = str(seed)
 np.random.seed(seed)
 torch.manual_seed(seed)
 torch.cuda.manual_seed(seed)
 torch.backends.cudnn.deterministic = True
seed_everything()<categorify>
class Generator(nn.Module):
def __init__(self, nz=128, channels=3):
super(Generator, self ).__init__()
self.nz = nz
self.channels = channels
def convlayer(n_input, n_output, k_size=4, stride=2, padding=0):
block = [
spectral_norm(nn.ConvTranspose2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False)) ,
nn.ReLU(inplace=True),
]
return block
self.layer1 = nn.Sequential(
 *convlayer(self.nz, 64, 4, 1, 0), 
)
 self.layer2 = nn.Sequential(ResBlockGenerator(64, 64, 2))
self.CBN2 = CategoricalConditionalBatchNorm2d(120,64)
 self.layer3 = nn.Sequential(ResBlockGenerator(64, 64, 2))
self.CBN3 = CategoricalConditionalBatchNorm2d(120,64)
 self.layer4 = nn.Sequential(ResBlockGenerator(64, 64, 2))
self.CBN4 = CategoricalConditionalBatchNorm2d(120,64)
 self.layer5 = nn.Sequential(ResBlockGenerator(64, 64, 2))
self.CBN5 = CategoricalConditionalBatchNorm2d(120,64)
 self.outlay = spectral_norm(nn.Conv2d(64, 3, 3, 1, 1))
self.Tanh = nn.Tanh()
def forward(self, z, y):
z = z.view(-1, self.nz, 1, 1)
 z = self.layer1(z)
 z = self.layer2(z)
 z = self.CBN2(z,y)
 z = self.layer3(z)
 z = self.CBN3(z,y)
 z = self.layer4(z)
 z = self.CBN4(z,y)
 z = self.layer5(z)
 z = self.CBN5(z,y)
 z = self.outlay(z)
 img = self.Tanh(z)
 return img
class Discriminator(nn.Module):
def __init__(self, channels=3):
super(Discriminator, self ).__init__()
self.channels = channels
def convlayer(n_input, n_output, k_size=4, stride=2, padding=0, bn=False):
block =[]
if bn:
block.append(spectral_norm(nn.Conv2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False)))
 block.append(nn.BatchNorm2d(n_output))
else:
block.append(nn.Conv2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False))
block.append(nn.LeakyReLU(0.2, inplace=True))
return block
self.model = nn.Sequential(
 *convlayer(self.channels, 64, 4, 2, 1),
ResBlockDiscriminator(64,128,2),
ResBlockDiscriminator(128,256,2),
ResBlockDiscriminator(256,728,2),
ResBlockDiscriminator(728,728,1), 
)
 self.logic_out = nn.Conv2d(728, 1, 4, 1, 0, bias=False)
 self.aux_out = nn.Conv2d(728, 120, 4, 1, 0, bias=False)
 def forward(self, imgs, label):
out = self.model(imgs)
 logic_output = self.logic_out(out)
 aux_output = self.aux_out(out)
 return logic_output.view(-1, 1),aux_output.view(-1,120 )
def truncated_normal(size, threshold=1):
values = truncnorm.rvs(-threshold, threshold, size=size)
 return values
def mse(imageA, imageB):
err = np.sum(( imageA.astype("float")- imageB.astype("float")) ** 2)
 err /= float(imageA.shape[0] * imageA.shape[1])
 return err
def make_answer() :
good_breeds = analyse_generated_by_class()
create_submit(good_breeds)
 def analyse_generated_by_class(n_images=5):
good_breeds = []
for l in range(len(decoded_dog_labels)) :
sample = []
for _ in range(n_images):
noise = torch.randn(1, nz, 1, 1, device=device)
 dog_label = torch.full(( 1,), l, device=device, dtype=torch.long)
 gen_image = netG(noise, dog_label ).to("cpu" ).clone().detach().squeeze(0)
 gen_image = gen_image.numpy().transpose(1, 2, 0)
 sample.append(gen_image)
 d = np.round(np.sum([mse(sample[k], sample[k+1])for k in range(len(sample)-1)])/n_images, 1)
 if d < 1.0: continue
print(f"Generated breed({d}): ", decoded_dog_labels[l])
 figure, axes = plt.subplots(1, len(sample), figsize=(64, 64))
for index, axis in enumerate(axes):
axis.axis('off')
 image_array =(sample[index] + 1.) / 2.
 axis.imshow(image_array)
 plt.show()
good_breeds.append(l)
 return good_breeds
def create_submit(good_breeds):
print("Creating submit")
 os.makedirs('.. /output_images', exist_ok=True)
 im_batch_size = 100
n_images = 10000
all_dog_labels = np.random.choice(good_breeds, size=n_images, replace=True)
 for i_batch in range(0, n_images, im_batch_size):
noise = torch.randn(im_batch_size, nz, 1, 1, device=device)
 dog_labels = torch.from_numpy(all_dog_labels[i_batch:(i_batch+im_batch_size)] ).to(device)
 gen_images = netG(noise, dog_labels)
 gen_images =(gen_images.to("cpu" ).clone().detach() + 1.) / 2.
 for ii, img in enumerate(gen_images):
save_image(gen_images[ii, :, :, :], os.path.join('.. /output_images', f'image_{i_batch + ii:05d}.png'))
shutil.make_archive('images', 'zip', '.. /output_images' )<load_pretrained>
class DogsDataset(Dataset):
def __init__(self, root, annotation_root, transform=None,
target_transform=None, loader=default_loader, n_process=4):
self.root = Path(root)
 self.annotation_root = Path(annotation_root)
 self.transform = transform
self.target_transform = target_transform
self.loader = loader
self.imgs= self.cut_out_dogs(n_process)
 def _get_annotation_path(self, img_path):
dog = Path(img_path ).stem
breed = dog.split('_')[0]
breed_dir = next(self.annotation_root.glob(f'{breed}-*'))
return breed_dir / dog
@staticmethod
def _get_dog_box(annotation_path):
tree = ET.parse(annotation_path)
 root = tree.getroot()
objects = root.findall('object')
 for o in objects:
bndbox = o.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)
 yield(xmin, ymin, xmax, ymax)
 def crop_dog(self, path):
imgs = []
annotation_path = self._get_annotation_path(path)
 dog = Path(path ).stem
label = label_number[dog.split('_')[0]]
for bndbox in self._get_dog_box(annotation_path):
img = self.loader(path)
 img_ = img.crop(bndbox)
 if np.sum(img_)!= 0:
img = img_
imgs.append([img,label])
 return imgs
def label_dog(self, path):
label = []
dog = Path(path ).stem
label.append(label_number[dog.split('_')[0]])
 return label
def cut_out_dogs(self, n_process):
with Pool(n_process)as p:
imgs = p.map(self.crop_dog, self.root.iterdir())
 return imgs
def __getitem__(self, index):
if self.transform is not None:
samples = self.transform(self.imgs[index][0][0])
 lables = self.imgs[index][0][1]
return samples,lables
def __len__(self):
return len(self.imgs )
class DataGenerator(Dataset):
def __init__(self, directory, transform=None, n_samples=np.inf, crop_dogs=True):
self.directory = directory
self.transform = transform
self.n_samples = n_samples
self.samples, self.labels = self.load_dogs_data(directory, crop_dogs)
 def load_dogs_data(self, directory, crop_dogs):
required_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(64),
torchvision.transforms.CenterCrop(64),
])
 imgs = []
labels = []
paths = []
for root, _, fnames in sorted(os.walk(directory)) :
for fname in sorted(fnames)[:min(self.n_samples, 999999999999999)]:
path = os.path.join(root, fname)
 paths.append(path)
 for path in paths:
try: img = dset.folder.default_loader(path)
 except: continue
annotation_basename = os.path.splitext(os.path.basename(path)) [0]
annotation_dirname = next(
 dirname for dirname in os.listdir('.. /input/annotation/Annotation/')if
dirname.startswith(annotation_basename.split('_')[0]))
if crop_dogs:
tree = ET.parse(os.path.join('.. /input/annotation/Annotation/',
annotation_dirname, annotation_basename))
root = tree.getroot()
objects = root.findall('object')
 for o in objects:
bndbox = o.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)
 object_img = required_transforms(img.crop(( xmin, ymin, xmax, ymax)))
 imgs.append(object_img)
 labels.append(annotation_dirname.split('-')[1].lower())
 else:
object_img = required_transforms(img)
 imgs.append(object_img)
 labels.append(annotation_dirname.split('-')[1].lower())
 return imgs, labels
def __getitem__(self, index):
sample = self.samples[index]
label = self.labels[index]
if self.transform is not None:
sample = self.transform(sample)
 return np.asarray(sample), label
def __len__(self):
return len(self.samples )<define_variables>
class ParamScheduler(object):
def __init__(self, optimizer, scale_fn, step_size):
if not isinstance(optimizer, Optimizer):
raise TypeError('{} is not an Optimizer'.format(
 type(optimizer ).__name__))
self.optimizer = optimizer
self.scale_fn = scale_fn
self.step_size = step_size
self.last_batch_iteration = 0
def batch_step(self):
for param_group in self.optimizer.param_groups:
param_group['lr'] = self.scale_fn(self.last_batch_iteration / self.step_size)
 self.last_batch_iteration += 1
def combine_scale_functions(scale_fns, phases=None):
if phases is None:
phases = [1./ len(scale_fns)] * len(scale_fns)
 phases = [phase / sum(phases)for phase in phases]
phases = torch.tensor([0] + phases)
 phases = torch.cumsum(phases, 0)
 def _inner(x):
idx =(x >= phases ).nonzero().max()
actual_x =(x - phases[idx])/(phases[idx + 1] - phases[idx])
 return scale_fns[idx](actual_x)
 return _inner
def scale_cos(start, end, x):
return start +(1 + np.cos(np.pi *(1 - x)))*(end - start)/ 2
%%time
database = '.. /input/all-dogs/all-dogs/'
random_transforms = [transforms.ColorJitter(brightness=(1,1.3), contrast=(1,1.3), saturation=0, hue=0), transforms.RandomRotation(degrees=5)]
transform = transforms.Compose([transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomApply(random_transforms, p=0.5),
transforms.ToTensor() ,
transforms.Normalize(( 0.5, 0.5, 0.5),(0.5, 0.5, 0.5)) ])
 train_data = DataGenerator(database, transform=transform, n_samples=25000, crop_dogs=True)
 decoded_dog_labels = {i:breed for i, breed in enumerate(sorted(set(train_data.labels)))}
encoded_dog_labels = {breed:i for i, breed in enumerate(sorted(set(train_data.labels)))}
train_data.labels = [encoded_dog_labels[l] for l in train_data.labels]
train_loader = torch.utils.data.DataLoader(train_data, shuffle=True,
batch_size=batch_size, num_workers=4 )<normalization>
random_transforms = [transforms.RandomRotation(degrees=5)]
transform = transforms.Compose([transforms.Resize(64),
transforms.CenterCrop(64),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomApply(random_transforms, p=0.3),
transforms.ToTensor() ,
transforms.Normalize(( 0.5, 0.5, 0.5),(0.5, 0.5, 0.5)) ])
 train_data = DogsDataset(TRAIN_DIR / 'all-dogs/', ANNOTATION_DIR, transform=transform)
 train_loader = DataLoader(train_data, shuffle=True, batch_size=batch_size, num_workers=4)
class PixelwiseNorm(nn.Module):
def __init__(self):
super(PixelwiseNorm, self ).__init__()
def forward(self, x, alpha=1e-8):
y = x.pow(2.).mean(dim=1, keepdim=True ).add(alpha ).sqrt()
y = x / y
return y
def show_generated_img_all(n_images=5, nz=128):
sample = []
for _ in range(n_images):
noise = torch.randn(1, nz, 1, 1, device=device)
 dog_label = torch.randint(0, len(encoded_dog_labels),(1,), device=device)
 gen_image = netG(noise, dog_label ).to("cpu" ).clone().detach().squeeze(0)
 gen_image = gen_image.numpy().transpose(1, 2, 0)
 sample.append(gen_image)
 figure, axes = plt.subplots(1, len(sample), figsize=(64, 64))
for index, axis in enumerate(axes):
axis.axis('off')
 image_array =(sample[index] + 1.) / 2.
 axis.imshow(image_array)
 plt.show()
def show_generated_img(n_images=5, nz=128):
sample = []
for _ in range(n_images):
noise = torch.randn(1, nz, 1, 1, device=device)
 dog_label = torch.randint(0, len(encoded_dog_labels),(1,), device=device)
 gen_image = netG(noise, dog_label ).to("cpu" ).clone().detach().squeeze(0)
 gen_image = gen_image.numpy().transpose(1, 2, 0)
 sample.append(gen_image)
 figure, axes = plt.subplots(1, len(sample), figsize=(64, 64))
for index, axis in enumerate(axes):
axis.axis('off')
 image_array =(sample[index] + 1.) / 2.
 axis.imshow(image_array)
 plt.show()
class MinibatchStdDev(nn.Module):
def __init__(self):
super(MinibatchStdDev, self ).__init__()
def forward(self, x, alpha=1e-8):
batch_size, _, height, width = x.shape
y = x - x.mean(dim=0, keepdim=True)
 y = torch.sqrt(y.pow(2.).mean(dim=0, keepdim=False)+ alpha)
 y = y.mean().view(1, 1, 1, 1)
 y = y.repeat(batch_size,1, height, width)
 y = torch.cat([x, y], 1)
 return y
def snconv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=False):
return spectral_norm(nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,
stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias))
def sndeconv2d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=False):
return spectral_norm(nn.ConvTranspose2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size,
stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias))
def snlinear(in_features, out_features):
return spectral_norm(nn.Linear(in_features=in_features, out_features=out_features))
def sn_embedding(num_embeddings, embedding_dim):
return spectral_norm(nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim))
class ConditionalBatchNorm2d(nn.Module):
def __init__(self, num_features, num_classes):
super().__init__()
self.num_features = num_features
self.bn = nn.BatchNorm2d(num_features, momentum=0.001, affine=False)
 self.embed = nn.Embedding(num_classes, num_features * 2)
 self.embed.weight.data[:, :num_features].fill_(1.)
self.embed.weight.data[:, num_features:].zero_()
def forward(self, x, y):
out = self.bn(x)
 gamma, beta = self.embed(y ).chunk(2, 1)
 out = gamma.view(-1, self.num_features, 1, 1)* out + beta.view(-1, self.num_features, 1, 1)
 return out<choose_model_class>
imgs,labels = next(iter(train_loader))
class GenBlock(nn.Module):
def __init__(self, in_channels, out_channels, num_classes):
super(GenBlock, self ).__init__()
self.cond_bn1 = ConditionalBatchNorm2d(in_channels, num_classes)
 self.relu = nn.ReLU(inplace=True)
 self.snconv2d1 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
 self.cond_bn2 = ConditionalBatchNorm2d(out_channels, num_classes)
 self.snconv2d2 = snconv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
 self.snconv2d0 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0)
 def forward(self, x, labels):
x0 = x
x = self.cond_bn1(x, labels)
 x = self.relu(x)
 x = F.interpolate(x, scale_factor=2, mode='nearest')
 x = self.snconv2d1(x)
 x = self.cond_bn2(x, labels)
 x = self.relu(x)
 x = self.snconv2d2(x)
 x0 = F.interpolate(x0, scale_factor=2, mode='nearest')
 x0 = self.snconv2d0(x0)
 out = x + x0
return out
class DiscBlock(nn.Module):
def __init__(self, in_channels, out_channels):
super(DiscBlock, self ).__init__()
self.relu = nn.ReLU()
self.snconv2d1 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
 self.snconv2d2 = snconv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1)
 self.downsample = nn.AvgPool2d(2)
 self.ch_mismatch = False
if in_channels != out_channels:
self.ch_mismatch = True
self.snconv2d0 = snconv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, stride=1, padding=0)
 def forward(self, x, downsample=True):
x0 = x
x = self.relu(x)
 x = self.snconv2d1(x)
 x = self.relu(x)
 x = self.snconv2d2(x)
 if downsample:
x = self.downsample(x)
 if downsample or self.ch_mismatch:
x0 = self.snconv2d0(x0)
 if downsample:
x0 = self.downsample(x0)
 out = x + x0
return out<define_search_model>
net_g = Generator(nz ).to(device)
 net_d = Discriminator().to(device)
 criterion = nn.BCELoss()
auxiliary_loss = torch.nn.CrossEntropyLoss()
scale_fn = combine_scale_functions(
 [partial(scale_cos, 1e-5, 5e-4), partial(scale_cos, 5e-4, 1e-4)], [0.2, 0.8])
 optimizer_g = optim.Adam(net_g.parameters() , lr=lr, betas=(beta1, 0.999))
optimizer_d = optim.Adam(net_d.parameters() , lr=lr, betas=(beta1, 0.999))
scheduler_g = ParamScheduler(optimizer_g, scale_fn, epochs * len(train_loader))
scheduler_d = ParamScheduler(optimizer_d, scale_fn, epochs * len(train_loader))
class Generator(nn.Module):
def __init__(self, nz, nfeats, nchannels, num_classes):
super(Generator, self ).__init__()
self.conv1 = spectral_norm(nn.ConvTranspose2d(nz, nfeats * 8, 4, 1, 0, bias=False))
self.block1 = GenBlock(nfeats * 8, nfeats * 8, num_classes)
 self.block2 = GenBlock(nfeats * 8, nfeats * 4, num_classes)
 self.block3 = GenBlock(nfeats * 4, nfeats * 2, num_classes)
 self.block4 = GenBlock(nfeats * 2, nfeats, num_classes)
 self.bn = nn.BatchNorm2d(nfeats, eps=1e-5, momentum=0.0001, affine=True)
 self.conv6 = spectral_norm(nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False))
self.pixnorm = PixelwiseNorm()
self.relu = nn.ReLU()
def forward(self, x, labels):
x = self.conv1(x)
 x = self.block1(x, labels)
 x = self.block2(x, labels)
 x = self.block3(x, labels)
 x = self.block4(x, labels)
 x = self.bn(x)
 x = self.relu(x)
 x = self.conv6(x)
 x = torch.tanh(x)
 return x
class Discriminator(nn.Module):
def __init__(self, nchannels, nfeats, num_classes):
super(Discriminator, self ).__init__()
self.label_emb = sn_embedding(num_classes, 64*64)
 self.conv1 = nn.Conv2d(nchannels+1, nfeats, 4, 2, 1, bias=False)
 self.block1 = DiscBlock(nfeats, nfeats*2)
 self.bn1 = nn.BatchNorm2d(nfeats * 2)
 self.block2 = DiscBlock(nfeats*2, nfeats*4)
 self.block3 = DiscBlock(nfeats*4, nfeats*8)
 self.block4 = DiscBlock(nfeats*8, nfeats*8)
 self.conv5 = spectral_norm(nn.Conv2d(nfeats * 8, 1, 2, 1, 0, bias=False))
def forward(self, imgs, labels):
enc = self.label_emb(labels ).view(( -1, 1, 64, 64))
enc = F.normalize(enc, p=2, dim=1)
 x = torch.cat(( imgs, enc), 1)
 x = self.conv1(x)
 x = self.block1(x)
 x = self.bn1(x)
 x = self.block2(x)
 x = self.block3(x)
 x = self.block4(x)
 x = F.leaky_relu(x, 0.2)
 x = torch.sigmoid(self.conv5(x))
return x.view(-1, 1)
<choose_model_class>
for epoch in range(epochs):
for i,(real_images,label)in enumerate(train_loader):
net_d.zero_grad()
real_images = real_images.to(device)
 label = label.type(torch.LongTensor ).to(device)
 batch_size = real_images.size(0)
 labels = torch.full(( batch_size, 1), real_label, device=device)
 scheduler_d.batch_step()
output_real,aux_real = net_d(real_images,label)
 noise = torch.randn(batch_size, nz, 1, 1, device=device)
 fake = net_g(noise,label)
 output_fake,aux_fake = net_d(fake.detach() ,label)
 loss_aux_real = auxiliary_loss(aux_real,label)
 loss_aux_fake = auxiliary_loss(aux_fake,label)
 err_d =(torch.mean(( output_real - torch.mean(output_fake)- labels)** 2)+
torch.mean(( output_fake - torch.mean(output_real)+ labels)** 2)) / 2
err_d = err_d*0.7 + loss_aux_real*0.15 + loss_aux_fake*0.15
err_d.backward(retain_graph=True)
 optimizer_d.step()
net_g.zero_grad()
scheduler_g.batch_step()
output_fake, aux_fake = net_d(fake,label)
 loss_aux_fake_g = auxiliary_loss(aux_fake,label)
 err_g =(torch.mean(( output_real - torch.mean(output_fake)+ labels)** 2)+
torch.mean(( output_fake - torch.mean(output_real)- labels)** 2)) / 2
err_g = err_g*0.85 + loss_aux_fake_g*0.15
err_g.backward()
optimizer_g.step()
print(f'[{epoch + 1}/{epochs}] Loss_d: {err_d.item() :.4f} Loss_g: {err_g.item() :.4f}' )
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 lr_d = 0.0006
lr_g = 0.0003
beta1 = 0.5
epochs = 300
num_classes = len(encoded_dog_labels)
 netG = Generator(128, 32, 3, num_classes ).to(device)
 netD = Discriminator(3, 48, num_classes ).to(device)
 criterion = nn.BCELoss()
criterionH = nn.MSELoss()
optimizerD = optim.Adam(netD.parameters() , lr=lr_d, betas=(beta1, 0.999))
optimizerG = optim.Adam(netG.parameters() , lr=lr_g, betas=(beta1, 0.999))
nz = 128
fixed_noise = torch.randn(25, nz, 1, 1, device=device)
 real_label = 0.7
fake_label = 0.0
batch_size = train_loader.batch_size
def get_model_weights(net):
average = {}
params = dict(net.named_parameters())
 for p in params:
average[p] = params[p].detach()
return average
print(sum(p.numel() for p in netG.parameters()))
print(sum(p.numel() for p in netD.parameters()))
averageD = False
averageG = False
hist_average_cutoff = 96 
<init_hyperparams>
def truncated_normal(size, threshold=1):
values = truncnorm.rvs(-threshold, threshold, size=size)
 return values
fids = {}
step = 1
for epoch in range(epochs):
if(time() - start)> 310 :
break
for ii,(real_images, dog_labels)in tqdm(enumerate(train_loader), total=len(train_loader)) :
if(time() - start)> 310 :
break
netD.zero_grad()
dog_labels = dog_labels.to(device)
 real_images = real_images.to(device)
 batch_size = real_images.size(0)
 labels = torch.full(( batch_size, 1), real_label, device=device)+ np.random.uniform(-0.1, 0.1)
 output = netD(real_images, dog_labels)
 errD_real = criterion(output, labels)
 errD_real.backward()
D_x = output.mean().item()
err_hD = 0
if epoch > hist_average_cutoff:
if not averageD:
print("Starting historical weight averaging for discriminator")
 averageD = get_model_weights(netD)
 paramsD = dict(netD.named_parameters())
 for p in paramsD:
err_hD += criterionH(paramsD[p], averageD[p])
 averageD[p] =(averageD[p] *(step-1)+ paramsD[p].detach())/step
err_hD.backward()
noise = torch.randn(batch_size, nz, 1, 1, device=device)
 fake = netG(noise, dog_labels)
 labels.fill_(fake_label)+ np.random.uniform(0, 0.2)
 output = netD(fake.detach() , dog_labels)
 errD_fake = criterion(output, labels)
 errD_fake.backward()
D_G_z1 = output.mean().item()
errD_final = errD_real + errD_fake + err_hD
optimizerD.step()
netG.zero_grad()
labels.fill_(real_label)
 output = netD(fake, dog_labels)
 errG = criterion(output, labels)
 errG.backward()
D_G_z2 = output.mean().item()
err_hG = 0
if epoch > hist_average_cutoff:
if not averageG:
print("Starting historical weight averaging for generator")
 averageG = get_model_weights(netG)
 paramsG = dict(netG.named_parameters())
 for p in paramsG:
err_hG += criterionH(paramsG[p], averageG[p])
 averageG[p] =(averageG[p] *(step-1)+ paramsG[p].detach())/step
err_hG.backward()
step += 1
errG_final = errG + err_hG
optimizerG.step()
if ii % 500 == 0:
print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)) : %.4f / %.4f'
%(epoch + 1, epochs, ii, len(train_loader),
errD_final.item() , errG_final.item() , D_x, D_G_z1, D_G_z2)) 
<load_pretrained>
im_batch_size = 50
n_images = 10000
for i_batch in range(0, n_images, im_batch_size):
z = truncated_normal(( im_batch_size, nz, 1, 1), threshold=1)
 gen_z = torch.from_numpy(z ).float().to(device)
 gen_labels = Variable(torch.LongTensor(np.random.randint(0, 120, im_batch_size)) ).to(device)
 gen_images =(net_g(gen_z,gen_labels)+ 1)/ 2
images = gen_images.to('cpu' ).clone().detach()
images = images.numpy().transpose(0, 2, 3, 1)
 for i_image in range(gen_images.size(0)) :
save_image(gen_images[i_image, :, :, :], OUT_DIR / f'image_{i_batch + i_image:05d}.png' )
Image.open(".. /input/all-dogs/all-dogs/n02110627_13662.jpg" )<import_modules>
shutil.make_archive('images', 'zip', OUT_DIR )