Kaggle competition: Generative Dog Images
<SOS> metric: PostProcessorKernel Kaggle data source: generative-dog-images<choose_model_class>

 %load_ext autoreload
%autoreload 2
%matplotlib inline
gc.enable()
torch.__version__
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.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 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.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<define_variables>
kernel_mode = True
if kernel_mode:
show_metric = False
show_animation = False
save_middle_model = False
if show_metric:
dataset_folder = ".. /input/generative-dog-images"
image_folder = ".. /input/generative-dog-images/all-dogs/all-dogs"
annotation_folder = ".. /input/generative-dog-images/annotation"
else:
dataset_folder = ".. /input"
image_folder = ".. /input/all-dogs"
annotation_folder = ".. /input/annotation"
else:
show_metric = True
show_animation = True
save_middle_model = False
dataset_folder = "."
image_folder = "all-dogs"
annotation_folder = "."
device = torch.device("cuda" if torch.cuda.is_available() else "cpu" )
train_dataiterator = get_dataiterator(all_images, all_labels, config['DataLoader'], device=device )<categorify>
batch_size = 32
max_training_time = 32_300
if kernel_mode:
max_iterations = 3_000_000
else:
max_iterations = 112_000
code_size = 128
n_mlp = 8
n_critic = 1
class Args:
base_lr = 0.0015
lr = {8: 0.002, 16: 0.004, 32: 0.006, 64: 0.008}
batch = {8: 128, 16: 64, 32: 32, 64: 32}
phase = {8: 800_000, 16: 400_000, 32: 400_000, 64: 400_000}
init_size = 8
max_size = 64
mixing = True
loss = 'r1'
args = Args()
beta1 = 0.0
beta2 = 0.99
transform = transforms.Compose([
transforms.RandomHorizontalFlip(p=0.3),
transforms.ToTensor() ,
transforms.Normalize(( 0.5, 0.5, 0.5),(0.5, 0.5, 0.5))
])
 if kernel_mode:
show_animation_iters = 2000
state_print_iters = 1000
else:
show_animation_iters = 2000
state_print_iters = 100
mifid_check_iters = 2000
mifid_check_images = 500
early_stopping_patience = 3
truncnorm_threshold = 1
if show_metric:
predict_batch_size = 25
else:
predict_batch_size = 50
netG = Generator(**config['Generator'] ).to(device, torch.float32)
 netD = Discriminator(**config['Discriminator'] ).to(device, torch.float32 )<load_pretrained>
start = time.time()
netGE = Generator(**config['Generator'] ).to(device, torch.float32)
 netGE.load_state_dict(netG.state_dict());<choose_model_class>
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
torch.backends.cudnn.benchmark = False
seed_everything()
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>
def load_dataset_images(root, n_samples=25000, image_size=64):
IMG_EXTENSIONS =('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif',
'.tiff', '.webp')
 def is_valid_file(x):
return datasets.folder.has_file_allowed_extension(x, IMG_EXTENSIONS)
 required_transforms = transforms.Compose([
transforms.Resize(image_size),
transforms.CenterCrop(image_size),
])
 imgs = []
paths = []
for root, _, fnames in sorted(os.walk(root)) :
for fname in sorted(fnames)[:min(n_samples, 999999999999999)]:
path = os.path.join(root, fname)
 paths.append(path)
 for path in paths:
if is_valid_file(path):
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(f'{annotation_folder}/Annotation/')
 if dirname.startswith(annotation_basename.split('_')[0]))
annotation_filename = os.path.join(
 f'{annotation_folder}/Annotation/', annotation_dirname,
annotation_basename)
 tree = ET.parse(annotation_filename)
 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)
 w = np.min(( xmax - xmin, ymax - ymin))
bbox =(xmin, ymin, xmin + w, ymin + w)
 object_img = required_transforms(img.crop(bbox))
imgs.append(object_img)
 print(f"Number of cropped dog images: {len(imgs)}")
 return imgs
in_memory_images = load_dataset_images(image_folder )
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>
class DataGenerator(torch.utils.data.Dataset):
def __init__(self,
dataset_images=None,
transform=None,
target_image_size=64):
self.dataset_images = dataset_images
self.transform = transform
self.target_image_size = target_image_size
self.samples = []
if self.target_image_size < 64:
print(f"Resizing images for DataGenerator to {target_image_size}x{target_image_size}...... ")
 resize_start_time = time.time()
required_transforms = transforms.Compose(
 [transforms.Resize(self.target_image_size)])
 for img in dataset_images:
self.samples.append(required_transforms(img))
print(f"Time spent on resizing images: " + \
f"{(time.time() - resize_start_time):.2f} seconds")
 else:
self.samples = self.dataset_images
def __getitem__(self, index):
sample = self.samples[index]
if self.transform is not None:
sample = self.transform(sample)
 return np.asarray(sample)
 def __len__(self):
return len(self.samples )
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>
def sample_data(batch_size, image_size=8):
train_data = DataGenerator(
 in_memory_images, transform=transform, target_image_size=image_size)
 train_loader = torch.utils.data.DataLoader(
 train_data, shuffle=True, batch_size=batch_size, num_workers=4)
 return train_loader
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>
criterion = nn.NLLLoss().to(device, torch.float32 )<define_variables>
def requires_grad(model, flag=True):
for p in model.parameters() :
p.requires_grad = flag
def accumulate(model1, model2, decay=0.999):
par1 = dict(model1.named_parameters())
 par2 = dict(model2.named_parameters())
 for k in par1.keys() :
par1[k].data.mul_(decay ).add_(1 - decay, par2[k].data)
 def adjust_lr(optimizer, lr):
for group in optimizer.param_groups:
mult = group.get('mult', 1)
 group['lr'] = lr * mult
def init_linear(linear):
init.xavier_normal(linear.weight)
 linear.bias.data.zero_()
def init_conv(conv, glu=True):
init.kaiming_normal(conv.weight)
 if conv.bias is not None:
conv.bias.data.zero_()
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>
class EqualLR:
def __init__(self, name):
self.name = name
def compute_weight(self, module):
weight = getattr(module, self.name + '_orig')
 fan_in = weight.data.size(1)* weight.data[0][0].numel()
return weight * sqrt(2 / fan_in)
 @staticmethod
def apply(module, name):
fn = EqualLR(name)
 weight = getattr(module, name)
 del module._parameters[name]
module.register_parameter(name + '_orig', nn.Parameter(weight.data))
module.register_forward_pre_hook(fn)
 return fn
def __call__(self, module, input):
weight = self.compute_weight(module)
 setattr(module, self.name, weight)
 def equal_lr(module, name='weight'):
EqualLR.apply(module, name)
 return module
class FusedUpsample(nn.Module):
def __init__(self, in_channel, out_channel, kernel_size, padding=0):
super().__init__()
weight = torch.randn(in_channel, out_channel, kernel_size, kernel_size)
 bias = torch.zeros(out_channel)
 fan_in = in_channel * kernel_size * kernel_size
self.multiplier = sqrt(2 / fan_in)
 self.weight = nn.Parameter(weight)
 self.bias = nn.Parameter(bias)
 self.pad = padding
def forward(self, input):
weight = F.pad(self.weight * self.multiplier, [1, 1, 1, 1])
 weight =(weight[:, :, 1:, 1:] + weight[:, :, :-1, 1:] +
weight[:, :, 1:, :-1] + weight[:, :, :-1, :-1])/ 4
out = F.conv_transpose2d(input,
weight,
self.bias,
stride=2,
padding=self.pad)
 return out
class FusedDownsample(nn.Module):
def __init__(self, in_channel, out_channel, kernel_size, padding=0):
super().__init__()
weight = torch.randn(out_channel, in_channel, kernel_size, kernel_size)
 bias = torch.zeros(out_channel)
 fan_in = in_channel * kernel_size * kernel_size
self.multiplier = sqrt(2 / fan_in)
 self.weight = nn.Parameter(weight)
 self.bias = nn.Parameter(bias)
 self.pad = padding
def forward(self, input):
weight = F.pad(self.weight * self.multiplier, [1, 1, 1, 1])
 weight =(weight[:, :, 1:, 1:] + weight[:, :, :-1, 1:] +
weight[:, :, 1:, :-1] + weight[:, :, :-1, :-1])/ 4
out = F.conv2d(input, weight, self.bias, stride=2, padding=self.pad)
 return out
class PixelNorm(nn.Module):
def __init__(self):
super().__init__()
def forward(self, input):
return input / torch.sqrt(
 torch.mean(input**2, dim=1, keepdim=True)+ 1e-8)
 class BlurFunctionBackward(Function):
@staticmethod
def forward(ctx, grad_output, kernel, kernel_flip):
ctx.save_for_backward(kernel, kernel_flip)
 grad_input = F.conv2d(grad_output,
kernel_flip,
padding=1,
groups=grad_output.shape[1])
 return grad_input
@staticmethod
def backward(ctx, gradgrad_output):
kernel, kernel_flip = ctx.saved_tensors
grad_input = F.conv2d(gradgrad_output,
kernel,
padding=1,
groups=gradgrad_output.shape[1])
 return grad_input, None, None
class BlurFunction(Function):
@staticmethod
def forward(ctx, input, kernel, kernel_flip):
ctx.save_for_backward(kernel, kernel_flip)
 output = F.conv2d(input, kernel, padding=1, groups=input.shape[1])
 return output
@staticmethod
def backward(ctx, grad_output):
kernel, kernel_flip = ctx.saved_tensors
grad_input = BlurFunctionBackward.apply(grad_output, kernel,
kernel_flip)
 return grad_input, None, None
blur = BlurFunction.apply
class Blur(nn.Module):
def __init__(self, channel):
super().__init__()
weight = torch.tensor([[1, 2, 1], [2, 4, 2], [1, 2, 1]],
dtype=torch.float32)
 weight = weight.view(1, 1, 3, 3)
 weight = weight / weight.sum()
weight_flip = torch.flip(weight, [2, 3])
 self.register_buffer('weight', weight.repeat(channel, 1, 1, 1))
self.register_buffer('weight_flip',
weight_flip.repeat(channel, 1, 1, 1))
def forward(self, input):
return blur(input, self.weight, self.weight_flip)
 class EqualConv2d(nn.Module):
def __init__(self, *args, **kwargs):
super().__init__()
conv = nn.Conv2d(*args, **kwargs)
 conv.weight.data.normal_()
conv.bias.data.zero_()
conv = conv
self.conv = equal_lr(conv)
 def forward(self, input):
return self.conv(input)
 class EqualLinear(nn.Module):
def __init__(self, in_dim, out_dim):
super().__init__()
linear = nn.Linear(in_dim, out_dim)
 linear.weight.data.normal_()
linear.bias.data.zero_()
self.linear = equal_lr(linear)
 def forward(self, input):
return self.linear(input)
 class ConvBlock(nn.Module):
def __init__(
 self,
in_channel,
out_channel,
kernel_size,
padding,
kernel_size2=None,
padding2=None,
downsample=False,
fused=False, 
):
super().__init__()
pad1 = padding
pad2 = padding
if padding2 is not None:
pad2 = padding2
kernel1 = kernel_size
kernel2 = kernel_size
if kernel_size2 is not None:
kernel2 = kernel_size2
self.conv1 = nn.Sequential(
 EqualConv2d(in_channel, out_channel, kernel1, padding=pad1),
nn.LeakyReLU(0.2), 
)
 if downsample:
if fused:
self.conv2 = nn.Sequential(
 Blur(out_channel),
FusedDownsample(out_channel,
out_channel,
kernel2,
padding=pad2),
nn.LeakyReLU(0.2), 
)
 else:
self.conv2 = nn.Sequential(
 Blur(out_channel),
EqualConv2d(out_channel,
out_channel,
kernel2,
padding=pad2),
nn.AvgPool2d(2),
nn.LeakyReLU(0.2), 
)
 else:
self.conv2 = nn.Sequential(
 EqualConv2d(out_channel, out_channel, kernel2, padding=pad2),
nn.LeakyReLU(0.2), 
)
 def forward(self, input):
out = self.conv1(input)
 out = self.conv2(out)
 return out
class AdaptiveInstanceNorm(nn.Module):
def __init__(self, in_channel, style_dim):
super().__init__()
self.norm = nn.InstanceNorm2d(in_channel)
 self.style = EqualLinear(style_dim, in_channel * 2)
 self.style.linear.bias.data[:
in_channel] = 1
self.style.linear.bias.data[in_channel:] = 0
def forward(self, input, style):
style = self.style(style ).unsqueeze(2 ).unsqueeze(3)
 gamma, beta = style.chunk(2, 1)
 out = self.norm(input)
 out = gamma * out + beta
return out
class NoiseInjection(nn.Module):
def __init__(self, channel):
super().__init__()
self.weight = nn.Parameter(torch.zeros(1, channel, 1, 1))
def forward(self, image, noise):
return image + self.weight * noise
class ConstantInput(nn.Module):
def __init__(self, channel, size=4):
super().__init__()
self.input = nn.Parameter(torch.randn(1, channel, size, size))
def forward(self, input):
batch = input.shape[0]
out = self.input.repeat(batch, 1, 1, 1)
 return out
class StyledConvBlock(nn.Module):
def __init__(
 self,
in_channel,
out_channel,
kernel_size=3,
padding=1,
style_dim=code_size,
initial=False,
upsample=False,
fused=False, 
):
super().__init__()
if initial:
self.conv1 = ConstantInput(in_channel)
 else:
if upsample:
if fused:
self.conv1 = nn.Sequential(
 FusedUpsample(in_channel,
out_channel,
kernel_size,
padding=padding),
Blur(out_channel), 
)
 else:
self.conv1 = nn.Sequential(
 nn.Upsample(scale_factor=2, mode='nearest'),
EqualConv2d(in_channel,
out_channel,
kernel_size,
padding=padding),
Blur(out_channel), 
)
 else:
self.conv1 = EqualConv2d(in_channel,
out_channel,
kernel_size,
padding=padding)
 self.noise1 = equal_lr(NoiseInjection(out_channel))
self.adain1 = AdaptiveInstanceNorm(out_channel, style_dim)
 self.lrelu1 = nn.LeakyReLU(0.2)
 self.conv2 = EqualConv2d(out_channel,
out_channel,
kernel_size,
padding=padding)
 self.noise2 = equal_lr(NoiseInjection(out_channel))
self.adain2 = AdaptiveInstanceNorm(out_channel, style_dim)
 self.lrelu2 = nn.LeakyReLU(0.2)
 def forward(self, input, style, noise):
out = self.conv1(input)
 out = self.noise1(out, noise)
 out = self.lrelu1(out)
 out = self.adain1(out, style)
 out = self.conv2(out)
 out = self.noise2(out, noise)
 out = self.lrelu2(out)
 out = self.adain2(out, style)
 return out
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>
class Generator(nn.Module):
def __init__(self, code_dim, fused=True):
super().__init__()
self.progression = nn.ModuleList([
StyledConvBlock(128, 128, 3, 1, initial=True),
StyledConvBlock(128, 128, 3, 1, upsample=True),
StyledConvBlock(128, 64, 3, 1, upsample=True),
StyledConvBlock(64, 32, 3, 1, upsample=True, fused=fused),
StyledConvBlock(32, 16, 3, 1, upsample=True, fused=fused),
])
 self.to_rgb = nn.ModuleList([
EqualConv2d(128, 3, 1),
EqualConv2d(128, 3, 1),
EqualConv2d(64, 3, 1),
EqualConv2d(32, 3, 1),
EqualConv2d(16, 3, 1),
])
 def forward(self, style, noise, step=0, alpha=-1, mixing_range=(-1, -1)) :
out = noise[0]
if len(style)< 2:
inject_index = [len(self.progression)+ 1]
else:
inject_index = random.sample(list(range(step)) , len(style)- 1)
 crossover = 0
for i,(conv, to_rgb)in enumerate(zip(self.progression, self.to_rgb)) :
if mixing_range ==(-1, -1):
if crossover < len(
 inject_index)and i > inject_index[crossover]:
crossover = min(crossover + 1, len(style))
style_step = style[crossover]
else:
if mixing_range[0] <= i <= mixing_range[1]:
style_step = style[1]
else:
style_step = style[0]
if i > 0 and step > 0:
out_prev = out
out = conv(out, style_step, noise[i])
 else:
out = conv(out, style_step, noise[i])
 if i == step:
out = to_rgb(out)
 if i > 0 and 0 <= alpha < 1:
skip_rgb = self.to_rgb[i - 1](out_prev)
 skip_rgb = F.interpolate(skip_rgb,
scale_factor=2,
mode='nearest')
 out =(1 - alpha)* skip_rgb + alpha * out
break
return out
class StyledGenerator(nn.Module):
def __init__(self, code_dim=128, n_mlp=8):
super().__init__()
self.generator = Generator(code_dim)
 layers = [PixelNorm() ]
for i in range(n_mlp):
layers.append(EqualLinear(code_dim, code_dim))
layers.append(nn.LeakyReLU(0.2))
self.style = nn.Sequential(*layers)
 def forward(
 self,
input,
noise=None,
step=0,
alpha=-1,
mean_style=None,
style_weight=0,
mixing_range=(-1, -1), 
):
styles = []
if type(input)not in(list, tuple):
input = [input]
for i in input:
styles.append(self.style(i))
batch = input[0].shape[0]
if noise is None:
noise = []
for i in range(step + 1):
size = 4 * 2**i
noise.append(
 torch.randn(batch, 1, size, size, device=input[0].device))
if mean_style is not None:
styles_norm = []
for style in styles:
styles_norm.append(mean_style + style_weight * 
(style - mean_style))
styles = styles_norm
return self.generator(styles,
noise,
step,
alpha,
mixing_range=mixing_range)
 def mean_style(self, input):
style = self.style(input ).mean(0, keepdim=True)
 return style
def truncated_normal(size, threshold=2.0, dtype=torch.float32, device='cpu'):
x = scipy.stats.truncnorm.rvs(-threshold, threshold, size=size)
 x = torch.from_numpy(x ).to(device, dtype)
 return x<define_variables>
class Discriminator(nn.Module):
def __init__(self, fused=True):
super().__init__()
self.progression = nn.ModuleList([
ConvBlock(16, 32, 3, 1, downsample=True, fused=fused),
ConvBlock(32, 64, 3, 1, downsample=True, fused=fused),
ConvBlock(64, 128, 3, 1, downsample=True),
ConvBlock(128, 128, 3, 1, downsample=True),
ConvBlock(129, 128, 3, 1, 4, 0),
])
 self.from_rgb = nn.ModuleList([
EqualConv2d(3, 16, 1),
EqualConv2d(3, 32, 1),
EqualConv2d(3, 64, 1),
EqualConv2d(3, 128, 1),
EqualConv2d(3, 128, 1),
])
 self.n_layer = len(self.progression)
 self.linear = EqualLinear(code_size, 1)
 def forward(self, input, step=0, alpha=-1):
for i in range(step, -1, -1):
index = self.n_layer - i - 1
if i == step:
out = self.from_rgb[index](input)
 if i == 0:
out_std = torch.sqrt(out.var(0, unbiased=False)+ 1e-8)
 mean_std = out_std.mean()
mean_std = mean_std.expand(out.size(0), 1, 4, 4)
 out = torch.cat([out, mean_std], 1)
 out = self.progression[index](out)
 if i > 0:
if i == step and 0 <= alpha < 1:
skip_rgb = F.avg_pool2d(input, 2)
 skip_rgb = self.from_rgb[index + 1](skip_rgb)
 out =(1 - alpha)* skip_rgb + alpha * out
out = out.squeeze(2 ).squeeze(2)
 out = self.linear(out)
 return out
def generate_eval_samples(generator, batch_size, latent_dim, num_classes):
latents = truncated_normal(( batch_size, latent_dim), dtype=torch.float32, device=device)
 labels = torch.randint(0, num_classes, size=(batch_size,), dtype=torch.long, device=device)
 with torch.no_grad() :
imgs =(generator(latents, labels)+ 1)/ 2
return imgs<import_modules>
import numpy as np
import os
import gzip, pickle
import tensorflow as tf
from scipy import linalg
import pathlib
import urllib
import warnings
from tqdm import tqdm
from PIL import Image
import shutil
import os
import random
from time import time
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.nn.utils import spectral_norm
from torch.utils.data import Dataset
import torchvision
import xml.etree.ElementTree as ET
from tqdm import tqdm_notebook as tqdm
from scipy.stats import truncnorm 
<define_variables>
class KernelEvalException(Exception):
pass
def create_model_graph(pth):

 with tf.gfile.GFile(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=32, 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 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'])
 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))
mean_min_d = np.mean(np.min(d, axis=1))
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
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=32,
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, n_images,
is_checksize, is_check_png):
path = pathlib.Path(path)
 files = list(path.glob('*.jpg')) + list(path.glob('*.png'))
seed_everything()
np.random.shuffle(files)
 files = files[:n_images]
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
gc.collect()
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,
n_images=1000,
feature_path=None):

 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,
n_images=n_images,
is_checksize=True,
is_check_png=True)
 if feature_path is None:
m2, s2, features2 = _handle_path_memorization(paths[1],
sess,
model_name,
n_images=n_images,
is_checksize=False,
is_check_png=False)
 else:
with np.load(feature_path)as f:
m2, s2, features2 = f['m'], f['s'], f['features']
fid_value = calculate_frechet_distance(m1, s1, m2, s2)
 distance = cosine_distance(features1, features2)
 return fid_value, distance
CHANNEL = {'RGB': 3, 'BW': 1}
IMG_CHANNEL = CHANNEL['RGB']
DATA_LOCATION = '.. /input/all-dogs/all-dogs/'
BATCH_SIZE = 32
LATENT_DIM = 128
CONV_DEPTHS_G = 32
CONV_DEPTHS_D = 48
LEARNING_RATE_G = 0.0003
LEARNING_RATE_D = 0.0001
BETA_1 = 0.5
T0_interval = 200
EPOCHS = 600
NUM_CLASSES = 10
TIME_FOR_TRAIN = 25000 
<set_options>
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
}
}
user_images_unzipped_path = '.. /output_images'
images_path = [user_images_unzipped_path, image_folder]
if kernel_mode:
public_path = '.. /input/dog-face-generation-competition-kid-metric-input/classify_image_graph_def.pb'
else:
public_path = 'classify_image_graph_def.pb'
fid_epsilon = 10e-15
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()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
<load_pretrained>
def compute_MiFID(n_images=1000):
seed_everything()
tf.set_random_seed(42)
 fid_value_public, distance_public = calculate_kid_given_paths(
 images_path, 'Inception', public_path, n_images)
 distance_public = distance_thresholding(
 distance_public, model_params['Inception']['cosine_distance_eps'])
 miFID = fid_value_public /(distance_public + fid_epsilon)
 print(
 f"FID_public: {fid_value_public}, distance_public: {distance_public} "
+ f"multiplied_public: {miFID}")
 return miFID
class DataGenerator(Dataset):
def __init__(self, directory, transform=None, n_samples=np.inf):
self.directory = directory
self.transform = transform
self.n_samples = n_samples
self.samples = self._load_subfolders_images(directory)
 if len(self.samples)== 0:
raise RuntimeError("Found 0 files in subfolders of: {}".format(directory))
def _load_subfolders_images(self, root):
IMG_EXTENSIONS =(
 '.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')
 def is_valid_file(x):
return torchvision.datasets.folder.has_file_allowed_extension(x, IMG_EXTENSIONS)
 required_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(64),
torchvision.transforms.CenterCrop(64),
])
 imgs = []
paths = []
for root, _, fnames in sorted(os.walk(root)) :
for fname in sorted(fnames)[:min(self.n_samples, 999999999999999)]:
path = os.path.join(root, fname)
 paths.append(path)
 for path in paths:
if is_valid_file(path):
img = torchvision.datasets.folder.default_loader(path)
 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]))
annotation_filename = os.path.join('.. /input/annotation/Annotation/',
annotation_dirname, annotation_basename)
 tree = ET.parse(annotation_filename)
 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)
 w = np.min(( xmax - xmin, ymax - ymin))
bbox =(xmin, ymin, xmin + w, ymin + w)
 object_img = required_transforms(img.crop(bbox))
imgs.append(object_img)
 return imgs
def __getitem__(self, index):
sample = self.samples[index]
if self.transform is not None:
sample = self.transform(sample)
 return np.asarray(sample)
 def __len__(self):
return len(self.samples)
 transform = torchvision.transforms.Compose([torchvision.transforms.RandomHorizontalFlip(p=0.1),
torchvision.transforms.ToTensor() ,
torchvision.transforms.Normalize(( 0.5, 0.5, 0.5),(0.5, 0.5, 0.5)) ])
 data_gen = DataGenerator(DATA_LOCATION, transform=transform, n_samples=25000)
 train_loader = torch.utils.data.DataLoader(data_gen, shuffle=True, batch_size=BATCH_SIZE, num_workers=4 )<normalization>
generator = StyledGenerator(code_size, n_mlp ).cuda()
discriminator = Discriminator().cuda()
g_running = StyledGenerator(code_size, n_mlp ).cuda()
g_running.train(False )
class PixelwiseNorm(torch.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 Generator(torch.nn.Module):
def __init__(self, latent_dim, nfeats, nchannels):
super(Generator, self ).__init__()
self.conv1 = spectral_norm(torch.nn.ConvTranspose2d(latent_dim, nfeats * 8, 4, 1, 0, bias=False))
self.conv2 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 8, nfeats * 8, 4, 2, 1, bias=False))
self.conv3 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 8, nfeats * 4, 4, 2, 1, bias=False))
self.conv4 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 4, nfeats * 2, 4, 2, 1, bias=False))
self.conv5 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 2, nfeats, 4, 2, 1, bias=False))
self.conv6 = spectral_norm(torch.nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False))
self.pixnorm = PixelwiseNorm()
def forward(self, x):
x = torch.nn.functional.leaky_relu(self.conv1(x))
x = self.pixnorm(x)
 x = torch.nn.functional.leaky_relu(self.conv2(x))
x = torch.nn.Dropout(0.05 )(x)
 x = self.pixnorm(x)
 x = torch.nn.functional.leaky_relu(self.conv3(x))
x = self.pixnorm(x)
 x = torch.nn.functional.leaky_relu(self.conv4(x))
x = torch.nn.Dropout(0.05 )(x)
 x = self.pixnorm(x)
 x = torch.nn.functional.leaky_relu(self.conv5(x))
x = self.pixnorm(x)
 x = torch.tanh(self.conv6(x))
return x 
<categorify>
g_optimizer = optim.Adam(generator.generator.parameters() ,
lr=args.base_lr,
betas=(beta1, beta2))
g_optimizer.add_param_group({
'params': generator.style.parameters() ,
'lr': args.base_lr * 0.01,
'mult': 0.01,
})
 d_optimizer = optim.Adam(discriminator.parameters() ,
lr=args.base_lr,
betas=(beta1, beta2))
accumulate(g_running, generator, 0 )
class MinibatchStdDev(torch.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<define_search_model>
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
 def save_models(model_name):
torch.save(
 {
'generator': generator.state_dict() ,
'discriminator': discriminator.state_dict() ,
'g_optimizer': g_optimizer.state_dict() ,
'd_optimizer': d_optimizer.state_dict() ,
'g_running': g_running.state_dict()
},
f'{model_name}.model', 
)
 print(
 f"Generator Model Size: {count_parameters(generator):,} trainable parameters" 
)
 print(
 f"Discriminator Model Size: {count_parameters(discriminator):,} trainable parameters"
)
class Discriminator(torch.nn.Module):
def __init__(self, nchannels, nfeats):
super(Discriminator, self ).__init__()
self.conv1 = torch.nn.Conv2d(nchannels, nfeats, 4, 2, 1, bias=False)
 self.conv2 = spectral_norm(torch.nn.Conv2d(nfeats, nfeats * 2, 4, 2, 1, bias=False))
self.bn2 = torch.nn.BatchNorm2d(nfeats * 2)
 self.conv3 = spectral_norm(torch.nn.Conv2d(nfeats * 2, nfeats * 4, 4, 2, 1, bias=False))
self.bn3 = torch.nn.BatchNorm2d(nfeats * 4)
 self.conv4 = spectral_norm(torch.nn.Conv2d(nfeats * 4, nfeats * 8, 4, 2, 1, bias=False))
self.bn4 = torch.nn.MaxPool2d(2)
 self.batch_discriminator = MinibatchStdDev()
self.conv5 = spectral_norm(torch.nn.Conv2d(nfeats * 8 + 1, 1, 2, 1, 0, bias=False))
def forward(self, x):
x = torch.nn.functional.leaky_relu(self.conv1(x), 0.2)
 x = torch.nn.functional.leaky_relu(self.bn2(self.conv2(x)) , 0.2)
 x = torch.nn.functional.leaky_relu(self.bn3(self.conv3(x)) , 0.2)
 x = torch.nn.functional.leaky_relu(self.bn4(self.conv4(x)) , 0.2)
 x = self.batch_discriminator(x)
 x = torch.sigmoid(self.conv5(x))
return x.view(-1, 1)
<compute_test_metric>
torch.cuda.empty_cache()
def truncated_normal(size, threshold=1):
values = truncnorm.rvs(-threshold, threshold, size=size)
 return values<train_model>
early_stopper = EarlyStoppingCriterion(patience=early_stopping_patience,
mode='min',
min_delta=1e-4)
 G_losses = []
D_losses = []
disc_loss_val = 0
gen_loss_val = 0
grad_loss_val = 0
gc.collect()
print("Starting Training Loop...")
 train_start_time = time.time()
alpha = 0
used_sample = 0
step = int(math.log2(args.init_size)) - 2
resolution = 4 * 2**step
phase = args.phase.get(resolution, 400_000)
 loader = sample_data(args.batch.get(resolution, batch_size),
image_size=resolution)
 data_loader = iter(loader)
 max_step = int(math.log2(args.max_size)) - 2
final_progress = False
adjust_lr(g_optimizer, args.lr.get(resolution, args.base_lr))
adjust_lr(d_optimizer, args.lr.get(resolution, args.base_lr))
requires_grad(generator, False)
 requires_grad(discriminator, True)
 print(f"Training with resolution {resolution}x{resolution}(" + \
f"step: {step}, batch_size: {args.batch.get(resolution, batch_size)}, " + \
f"Generator LR: {g_optimizer.state_dict() ['param_groups'][0]['lr']}, " + \
f"Style LR: {g_optimizer.state_dict() ['param_groups'][1]['lr']}, " + \
f"Discriminator LR: {d_optimizer.state_dict() ['param_groups'][0]['lr']})")
 go_ahead = True
for iters in range(max_iterations):
if not go_ahead:
break
gc.collect()
iter_start_time = time.time()
end = time.time()
if(end - start)> max_training_time:
print(
 f"Trained over {max_training_time:,} seconds, stopped at iteration {iters+1}" 
)
 go_ahead = False
break
discriminator.zero_grad()
alpha = min(1, 1 / phase *(used_sample + 1))
if resolution == args.init_size or final_progress:
alpha = 1
if used_sample > phase * 2:
used_sample = 0
step += 1
if step > max_step:
step = max_step
final_progress = True
else:
alpha = 0
resolution = 4 * 2**step
phase = args.phase.get(resolution, 400_000)
 del loader
gc.collect()
loader = sample_data(args.batch.get(resolution, batch_size),
image_size=resolution)
 data_loader = iter(loader)
 adjust_lr(g_optimizer, args.lr.get(resolution, args.base_lr))
adjust_lr(d_optimizer, args.lr.get(resolution, args.base_lr))
print(f"Training with resolution {resolution}x{resolution}(" + \
f"step: {step}, batch_size: {args.batch.get(resolution, batch_size)}, " + \
f"Generator LR: {g_optimizer.state_dict() ['param_groups'][0]['lr']}, " + \
f"Style LR: {g_optimizer.state_dict() ['param_groups'][1]['lr']}, " + \
f"Discriminator LR: {d_optimizer.state_dict() ['param_groups'][0]['lr']})")
 try:
real_image = next(data_loader)
 except(OSError, StopIteration):
data_loader = iter(loader)
 real_image = next(data_loader)
 used_sample += real_image.shape[0]
b_size = real_image.size(0)
 real_image = real_image.cuda()
if args.loss == 'wgan-gp':
real_predict = discriminator(real_image, step=step, alpha=alpha)
 real_predict = real_predict.mean() - 0.001 *(real_predict**2 ).mean() 
(-real_predict ).backward()
elif args.loss == 'r1':
real_image.requires_grad = True
real_predict = discriminator(real_image, step=step, alpha=alpha)
 real_predict = F.softplus(-real_predict ).mean()
real_predict.backward(retain_graph=True)
 grad_real = grad(outputs=real_predict.sum() ,
inputs=real_image,
create_graph=True)[0]
grad_penalty =(grad_real.view(grad_real.size(0),
-1 ).norm(2, dim=1)**2 ).mean()
grad_penalty = 10 / 2 * grad_penalty
grad_penalty.backward()
grad_loss_val = grad_penalty.item()
if args.mixing and random.random() < 0.9:
gen_in11, gen_in12, gen_in21, gen_in22 = torch.randn(
 4, b_size, code_size, device='cuda' ).chunk(4, 0)
 gen_in1 = [gen_in11.squeeze(0), gen_in12.squeeze(0)]
gen_in2 = [gen_in21.squeeze(0), gen_in22.squeeze(0)]
else:
gen_in1, gen_in2 = torch.randn(2, b_size, code_size,
device='cuda' ).chunk(2, 0)
 gen_in1 = gen_in1.squeeze(0)
 gen_in2 = gen_in2.squeeze(0)
 fake_image = generator(gen_in1, step=step, alpha=alpha)
 fake_predict = discriminator(fake_image, step=step, alpha=alpha)
 if args.loss == 'wgan-gp':
fake_predict = fake_predict.mean()
fake_predict.backward()
eps = torch.rand(b_size, 1, 1, 1 ).cuda()
x_hat = eps * real_image.data +(1 - eps)* fake_image.data
x_hat.requires_grad = True
hat_predict = discriminator(x_hat, step=step, alpha=alpha)
 grad_x_hat = grad(outputs=hat_predict.sum() ,
inputs=x_hat,
create_graph=True)[0]
grad_penalty =(
(grad_x_hat.view(grad_x_hat.size(0), -1 ).norm(2, dim=1)-
1)**2 ).mean()
grad_penalty = 10 * grad_penalty
grad_penalty.backward()
grad_loss_val = grad_penalty.item()
disc_loss_val =(real_predict - fake_predict ).item()
elif args.loss == 'r1':
fake_predict = F.softplus(fake_predict ).mean()
fake_predict.backward()
disc_loss_val =(real_predict + fake_predict ).item()
d_optimizer.step()
if(iters + 1)% n_critic == 0:
generator.zero_grad()
requires_grad(generator, True)
 requires_grad(discriminator, False)
 fake_image = generator(gen_in2, step=step, alpha=alpha)
 predict = discriminator(fake_image, step=step, alpha=alpha)
 if args.loss == 'wgan-gp':
loss = -predict.mean()
elif args.loss == 'r1':
loss = F.softplus(-predict ).mean()
gen_loss_val = loss.item()
loss.backward()
g_optimizer.step()
accumulate(g_running, generator)
 requires_grad(generator, False)
 requires_grad(discriminator, True)
 if show_animation and(iters + 1)% show_animation_iters == 0:
show_generated_img(iters)
 if(iters + 1)% state_print_iters == 0:
print(f"[{iters+1}/{max_iterations}] Phase Size:{phase}, " + \
f"Resolution: {4 * 2 ** step}, " + \
f"Loss_G: {gen_loss_val:.4f}, " + \
f"Loss_D: {disc_loss_val:.4f}, " + \
f"Grad: {grad_loss_val:.4f}, Alpha: {alpha:.5f}, Used Samples: {used_sample:,}")
 print(f"Time spent on iteration {iters+1}: " + \
f"{(time.time() - iter_start_time):.2f} seconds")
 G_losses.append(gen_loss_val)
 D_losses.append(disc_loss_val)
 if show_metric and(
(iters + 1)%
mifid_check_iters)== 0 and resolution == args.max_size:
generate_validation_images(threshold=truncnorm_threshold,
n_images=mifid_check_images)
 val_mifid = compute_MiFID(mifid_check_images)
 if early_stopper.step(val_mifid):
if early_stopper.is_improved:
best_score = val_mifid
best_iters = iters
print(
 f'MiFID Improved: {val_mifid:.4f}, saving model...... ' 
)
 else:
print('No improvements in this iteration')
 gc.collect()
else:
print(
 f'No more improvement....(Best MiFID: {best_score:.6f}, {best_iters+1} iterations)' 
)
 gc.collect()
break
if show_metric:
print(f'Best MiFID: {best_score:.6f}, {best_iters+1} iterations')
 else:
best_iters = iters
print(f"Total time spent on training: " + \
f"{(time.time() - train_start_time)/60:.2f} minutes " + \
f"({best_iters + 1} iterations, step: {step}, alpha: {alpha})" )
def train_module(epochs):
step = 0
start = time()
for epoch in range(epochs):
for ii,(real_images)in enumerate(train_loader):
end = time()
if(end - start)> TIME_FOR_TRAIN:
break
netD.zero_grad()
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)
 errD_real = criterion(output, labels)
 errD_real.backward()
D_x = output.mean().item()
noise = torch.randn(batch_size, LATENT_DIM, 1, 1, device=device)
 fake = netG(noise)
 labels.fill_(fake_label)+ np.random.uniform(0, 0.2)
 output = netD(fake.detach())
 errD_fake = criterion(output, labels)
 errD_fake.backward()
D_G_z1 = output.mean().item()
errD = errD_real + errD_fake
optimizerD.step()
netG.zero_grad()
labels.fill_(real_label)
 output = netD(fake)
 errG = criterion(output, labels)
 errG.backward()
D_G_z2 = output.mean().item()
optimizerG.step()
if step % 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.item() , errG.item() , D_x, D_G_z1, D_G_z2))
valid_image = netG(fixed_noise)
 step += 1
lr_schedulerG.step(epoch)
 lr_schedulerD.step(epoch)
 if epoch % 5 == 0:
show_generated_img()
print(end - start)
 if(end - start)> TIME_FOR_TRAIN:
break<choose_model_class>
netG = Generator(LATENT_DIM, CONV_DEPTHS_G, IMG_CHANNEL ).to(device)
 netD = Discriminator(3, CONV_DEPTHS_D ).to(device)
 criterion = torch.nn.BCELoss()
optimizerD = torch.optim.Adam(netD.parameters() , lr=LEARNING_RATE_G, betas=(BETA_1, 0.999))
optimizerG = torch.optim.Adam(netG.parameters() , lr=LEARNING_RATE_D, betas=(BETA_1, 0.999))
lr_schedulerG = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizerG,
T_0=EPOCHS // T0_interval, eta_min=0.00005)
 lr_schedulerD = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizerD,
T_0=EPOCHS // T0_interval, eta_min=0.00005)
 fixed_noise = torch.randn(25, LATENT_DIM, 1, 1, device=device)
 real_label = 0.8
fake_label = 0.15
batch_size = train_loader.batch_size
train_module(EPOCHS )<define_variables>
if not show_metric:
final_model = f'stylegan_{best_iters+1}iters'
save_models(final_model)
 print(f"{final_model}.model saved." )
EPOCHS = 574
BATCH = 32
SIZE = 64
LRG = 0.0004
LRD = 0.0008
NZ = 64
FEATURES_G = 80
FEATURES_D = 112
LRFACTOR_G = 0.9
LRFACTOR_D = 0.9
REAL_LABEL = [0.9, 0.9]
FAKE_LABEL = [0.00, 0.00]
REAL_FAKE_LABEL = 0.9
SEED = 1234
RATIO = 1.50
PATH_ANNOTATIONS = '.. /input/annotation/Annotation'
PATH_ALL_DOGS = '.. /input/all-dogs/all-dogs'
PATH_DOGS = '.. /cropped-dogs' 
<set_options>
def truncated_normal(size, threshold=1):
values = truncnorm.rvs(-threshold, threshold, size=size)
 return values
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 def random_seed(seed):
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
random.seed(seed)
 np.random.seed(seed)
 torch.manual_seed(seed)
 tf.set_random_seed(seed)
 def R4(f): return F'{f:0.4}'.ljust(6, '0' ).ljust(8, ' ')
 def R6(f): return F'{f:0.6}'.ljust(8, '0' ).ljust(8, ' ')
 random_seed(SEED)
<import_modules>
%%time
if not os.path.exists('.. /output_images'):
os.mkdir('.. /output_images')
 else:
shutil.rmtree('.. /output_images', ignore_errors=True)
 os.mkdir('.. /output_images')
 im_batch_size = 50
n_images = 10000
first_batch = []
for i_batch in range(0, n_images, im_batch_size):
z = truncated_normal(( im_batch_size, code_size),
threshold=truncnorm_threshold)
 gen_z = torch.from_numpy(z ).float().to(device)
 with torch.no_grad() :
gen_images = g_running(gen_z, step=step, alpha=alpha ).data.cpu()
for i_image in range(gen_images.size(0)) :
new_image = gen_images[i_image, :, :, :]
if i_batch == 0:
first_batch.append(new_image)
 save_image(new_image,
os.path.join('.. /output_images',
f'image_{i_batch+i_image:05d}.png'),
normalize=True,
range=(-1, 1))
generated_image_count = len([name for name in os.listdir('.. /output_images')])
 print(f"Number of generated images: {generated_image_count}")
 shutil.make_archive('images', 'zip', '.. /output_images' )
<choose_model_class><EOS> 
%%time
show_generated_img_all(64 )