diff --git a/recognition/45857876/CustomLayers.py b/recognition/45857876/CustomLayers.py new file mode 100644 index 0000000000..29f4a7ddff --- /dev/null +++ b/recognition/45857876/CustomLayers.py @@ -0,0 +1,315 @@ + +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class PixelNormLayer(nn.Module): + def __init__(self, epsilon=1e-8): + super().__init__() + self.epsilon = epsilon + + def forward(self, x): + return x * torch.sqrt(torch.mean(x ** 2, dim=1, keepdim=True) + self.epsilon) + + +class Upscale2d(nn.Module): + @staticmethod + def upscale2d(x, factor=2, gain=1): + assert x.dim() == 4 + if gain != 1: + x = x * gain + if factor != 1: + shape = x.shape + x = x.view(shape[0], shape[1], shape[2], 1, shape[3], 1).expand(-1, -1, -1, factor, -1, factor) + x = x.contiguous().view(shape[0], shape[1], factor * shape[2], factor * shape[3]) + return x + + def __init__(self, factor=2, gain=1): + super().__init__() + assert isinstance(factor, int) and factor >= 1 + self.gain = gain + self.factor = factor + + def forward(self, x): + return self.upscale2d(x, factor=self.factor, gain=self.gain) + + +class Downscale2d(nn.Module): + def __init__(self, factor=2, gain=1): + super().__init__() + assert isinstance(factor, int) and factor >= 1 + self.gain = torch.tensor(gain) + self.factor = torch.tensor(factor) + if factor == 2: + f = [torch.sqrt(self.gain) / self.factor] * self.factor + self.blur = BlurLayer(kernel=f, normalize=False, stride=factor) + else: + self.blur = None + + def forward(self, x): + assert x.dim() == 4 + # 2x2, float32 => downscale using _blur2d(). + if self.blur is not None and x.dtype == torch.float32: + return self.blur(x) + + # Apply gain. + if self.gain != 1: + x = x * self.gain + + # No-op => early exit. + if self.factor == 1: + return x + + # Large factor => downscale using tf.nn.avg_pool(). + # NOTE: Requires tf_config['graph_options.place_pruned_graph']=True to work. + return F.avg_pool2d(x, self.factor) + + +class EqualizedLinear(nn.Module): + """Linear layer with equalized learning rate and custom learning rate multiplier.""" + + def __init__(self, input_size, output_size, gain=2 ** 0.5, use_wscale=False, lrmul=1, bias=True): + super().__init__() + he_std = gain * input_size ** (-0.5) # He init + # Equalized learning rate and custom learning rate multiplier. + if use_wscale: + init_std = 1.0 / lrmul + self.w_mul = he_std * lrmul + else: + init_std = he_std / lrmul + self.w_mul = lrmul + self.weight = torch.nn.Parameter(torch.randn(output_size, input_size) * init_std) + if bias: + self.bias = torch.nn.Parameter(torch.zeros(output_size)) + self.b_mul = lrmul + else: + self.bias = None + + def forward(self, x): + bias = self.bias + if bias is not None: + bias = bias * self.b_mul + return F.linear(x, self.weight * self.w_mul, bias) + + +class EqualizedConv2d(nn.Module): + """Conv layer with equalized learning rate and custom learning rate multiplier.""" + + def __init__(self, input_channels, output_channels, kernel_size, stride=1, gain=2 ** 0.5, use_wscale=False, + lrmul=1, bias=True, intermediate=None, upscale=False, downscale=False): + super().__init__() + if upscale: + self.upscale = Upscale2d() + else: + self.upscale = None + if downscale: + self.downscale = Downscale2d() + else: + self.downscale = None + he_std = gain * (input_channels * kernel_size ** 2) ** (-0.5) # He init + self.kernel_size = kernel_size + if use_wscale: + init_std = 1.0 / lrmul + self.w_mul = he_std * lrmul + else: + init_std = he_std / lrmul + self.w_mul = lrmul + self.weight = torch.nn.Parameter( + torch.randn(output_channels, input_channels, kernel_size, kernel_size) * init_std) + if bias: + self.bias = torch.nn.Parameter(torch.zeros(output_channels)) + self.b_mul = lrmul + else: + self.bias = None + self.intermediate = intermediate + + def forward(self, x): + bias = self.bias + if bias is not None: + bias = bias * self.b_mul + + have_convolution = False + if self.upscale is not None and min(x.shape[2:]) * 2 >= 128: + # this is the fused upscale + conv from StyleGAN, sadly this seems incompatible with the non-fused way + # this really needs to be cleaned up and go into the conv... + w = self.weight * self.w_mul + w = w.permute(1, 0, 2, 3) + # probably applying a conv on w would be more efficient. also this quadruples the weight (average)?! + w = F.pad(w, [1, 1, 1, 1]) + w = w[:, :, 1:, 1:] + w[:, :, :-1, 1:] + w[:, :, 1:, :-1] + w[:, :, :-1, :-1] + x = F.conv_transpose2d(x, w, stride=2, padding=(w.size(-1) - 1) // 2) + have_convolution = True + elif self.upscale is not None: + x = self.upscale(x) + + downscale = self.downscale + intermediate = self.intermediate + if downscale is not None and min(x.shape[2:]) >= 128: + w = self.weight * self.w_mul + w = F.pad(w, [1, 1, 1, 1]) + # in contrast to upscale, this is a mean... + w = (w[:, :, 1:, 1:] + w[:, :, :-1, 1:] + w[:, :, 1:, :-1] + w[:, :, :-1, :-1]) * 0.25 # avg_pool? + x = F.conv2d(x, w, stride=2, padding=(w.size(-1) - 1) // 2) + have_convolution = True + downscale = None + elif downscale is not None: + assert intermediate is None + intermediate = downscale + + if not have_convolution and intermediate is None: + return F.conv2d(x, self.weight * self.w_mul, bias, padding=self.kernel_size // 2) + elif not have_convolution: + x = F.conv2d(x, self.weight * self.w_mul, None, padding=self.kernel_size // 2) + + if intermediate is not None: + x = intermediate(x) + + if bias is not None: + x = x + bias.view(1, -1, 1, 1) + return x + + +class NoiseLayer(nn.Module): + """adds noise. noise is per pixel (constant over channels) with per-channel weight""" + + def __init__(self, channels): + super().__init__() + self.weight = nn.Parameter(torch.zeros(channels)) + self.noise = None + + def forward(self, x, noise=None): + if noise is None and self.noise is None: + noise = torch.randn(x.size(0), 1, x.size(2), x.size(3), device=x.device, dtype=x.dtype) + elif noise is None: + # here is a little trick: if you get all the noise layers and set each + # modules .noise attribute, you can have pre-defined noise. + # Very useful for analysis + noise = self.noise + x = x + self.weight.view(1, -1, 1, 1) * noise + return x + + +class StyleMod(nn.Module): + def __init__(self, latent_size, channels, use_wscale): + super(StyleMod, self).__init__() + self.lin = EqualizedLinear(latent_size, + channels * 2, + gain=1.0, use_wscale=use_wscale) + + def forward(self, x, latent): + style = self.lin(latent) # style => [batch_size, n_channels*2] + + shape = [-1, 2, x.size(1)] + (x.dim() - 2) * [1] + style = style.view(shape) # [batch_size, 2, n_channels, ...] + x = x * (style[:, 0] + 1.) + style[:, 1] + return x + + +class LayerEpilogue(nn.Module): + """Things to do at the end of each layer.""" + + def __init__(self, channels, dlatent_size, use_wscale, + use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer): + super().__init__() + + layers = [] + if use_noise: + layers.append(('noise', NoiseLayer(channels))) + layers.append(('activation', activation_layer)) + if use_pixel_norm: + layers.append(('pixel_norm', PixelNormLayer())) + if use_instance_norm: + layers.append(('instance_norm', nn.InstanceNorm2d(channels))) + + self.top_epi = nn.Sequential(OrderedDict(layers)) + + if use_styles: + self.style_mod = StyleMod(dlatent_size, channels, use_wscale=use_wscale) + else: + self.style_mod = None + + def forward(self, x, dlatents_in_slice=None): + x = self.top_epi(x) + if self.style_mod is not None: + x = self.style_mod(x, dlatents_in_slice) + else: + assert dlatents_in_slice is None + return x + + +class BlurLayer(nn.Module): + def __init__(self, kernel=None, normalize=True, flip=False, stride=1): + super(BlurLayer, self).__init__() + if kernel is None: + kernel = [1, 2, 1] + kernel = torch.tensor(kernel, dtype=torch.float32) + kernel = kernel[:, None] * kernel[None, :] + kernel = kernel[None, None] + if normalize: + kernel = kernel / kernel.sum() + if flip: + kernel = kernel[:, :, ::-1, ::-1] + self.register_buffer('kernel', kernel) + self.stride = stride + + def forward(self, x): + # expand kernel channels + kernel = self.kernel.expand(x.size(1), -1, -1, -1) + x = F.conv2d( + x, + kernel, + stride=self.stride, + padding=int((self.kernel.size(2) - 1) / 2), + groups=x.size(1) + ) + return x + + +class View(nn.Module): + def __init__(self, *shape): + super().__init__() + self.shape = shape + + def forward(self, x): + return x.view(x.size(0), *self.shape) + + +class StddevLayer(nn.Module): + def __init__(self, group_size=4, num_new_features=1): + super().__init__() + self.group_size = group_size + self.num_new_features = num_new_features + + def forward(self, x): + b, c, h, w = x.shape + group_size = min(self.group_size, b) + y = x.reshape([group_size, -1, self.num_new_features, + c // self.num_new_features, h, w]) + y = y - y.mean(0, keepdim=True) + y = (y ** 2).mean(0, keepdim=True) + y = (y + 1e-8) ** 0.5 + y = y.mean([3, 4, 5], keepdim=True).squeeze(3) # don't keep the meaned-out channels + y = y.expand(group_size, -1, -1, h, w).clone().reshape(b, self.num_new_features, h, w) + z = torch.cat([x, y], dim=1) + return z + + +class Truncation(nn.Module): + def __init__(self, avg_latent, max_layer=8, threshold=0.7, beta=0.995): + super().__init__() + self.max_layer = max_layer + self.threshold = threshold + self.beta = beta + self.register_buffer('avg_latent', avg_latent) + + def update(self, last_avg): + self.avg_latent.copy_(self.beta * self.avg_latent + (1. - self.beta) * last_avg) + + def forward(self, x): + assert x.dim() == 3 + interp = torch.lerp(self.avg_latent, x, self.threshold) + do_trunc = (torch.arange(x.size(1)) < self.max_layer).view(1, -1, 1).to(x.device) + return torch.where(do_trunc, interp, x) \ No newline at end of file diff --git a/recognition/45857876/generate_mixing_figure.py b/recognition/45857876/generate_mixing_figure.py new file mode 100644 index 0000000000..7f7df006a5 --- /dev/null +++ b/recognition/45857876/generate_mixing_figure.py @@ -0,0 +1,107 @@ +import os +import argparse +import numpy as np +from PIL import Image + +import torch + +from model import Generator + + +def adjust_dynamic_range(data, drange_in=(-1, 1), drange_out=(0, 1)): + """ + adjust the dynamic colour range of the given input data + """ + if drange_in != drange_out: + scale = (np.float32(drange_out[1]) - np.float32(drange_out[0])) / ( + np.float32(drange_in[1]) - np.float32(drange_in[0])) + bias = (np.float32(drange_out[0]) - np.float32(drange_in[0]) * scale) + data = data * scale + bias + return torch.clamp(data, min=0, max=1) + + +def draw_style_mixing_figure(png, gen, out_depth, src_seeds, dst_seeds, style_ranges): + n_col = len(src_seeds) + n_row = len(dst_seeds) + w = h = 2 ** (out_depth + 2) + with torch.no_grad(): + latent_size = gen.g_mapping.latent_size + # print(latent_size) + src_latents_np = np.stack([np.random.RandomState(seed).randn(latent_size, ) for seed in src_seeds]) + dst_latents_np = np.stack([np.random.RandomState(seed).randn(latent_size, ) for seed in dst_seeds]) + src_latents = torch.from_numpy(src_latents_np.astype(np.float32)) + dst_latents = torch.from_numpy(dst_latents_np.astype(np.float32)) + src_dlatents = gen.g_mapping(src_latents) # [seed, layer, component] + dst_dlatents = gen.g_mapping(dst_latents) # [seed, layer, component] + src_images = gen.g_synthesis(src_dlatents, depth=out_depth, alpha=1) + dst_images = gen.g_synthesis(dst_dlatents, depth=out_depth, alpha=1) + + src_dlatents_np = src_dlatents.numpy() + dst_dlatents_np = dst_dlatents.numpy() + canvas = Image.new('RGB', (w * (n_col + 1), h * (n_row + 1)), 'white') + for col, src_image in enumerate(list(src_images)): + src_image = adjust_dynamic_range(src_image) + src_image = src_image.mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy() + canvas.paste(Image.fromarray(src_image, 'RGB'), ((col + 1) * w, 0)) + for row, dst_image in enumerate(list(dst_images)): + dst_image = adjust_dynamic_range(dst_image) + dst_image = dst_image.mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy() + canvas.paste(Image.fromarray(dst_image, 'RGB'), (0, (row + 1) * h)) + + row_dlatents = np.stack([dst_dlatents_np[row]] * n_col) + row_dlatents[:, style_ranges[row]] = src_dlatents_np[:, style_ranges[row]] + row_dlatents = torch.from_numpy(row_dlatents) + + row_images = gen.g_synthesis(row_dlatents, depth=out_depth, alpha=1) + for col, image in enumerate(list(row_images)): + image = adjust_dynamic_range(image) + image = image.mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy() + canvas.paste(Image.fromarray(image, 'RGB'), ((col + 1) * w, (row + 1) * h)) + canvas.save(png) + + +def main(args): + """ + Main function for the generate mixing image + """ + + print("Creating generator object ...") + # create the generator object + gen = Generator(resolution=256, + num_channels=3, + structure="linear" + ) + + print("Loading the generator weights from:", args.generator_file) + # load the weights into it + # print(gen.load_state_dict) + gen.load_state_dict(torch.load(args.generator_file)) + + # path for saving the files: + # generate the images: + # src_seeds = [639, 701, 687, 615, 1999], dst_seeds = [888, 888, 888], + draw_style_mixing_figure(os.path.join('figure03-style-mixing.png'), gen, + out_depth=6, src_seeds=[639, 1995, 687, 615, 1999], dst_seeds=[888, 888, 888], + style_ranges=[range(0, 2)] * 1 + [range(2, 8)] * 1 + [range(8, 14)] * 1) + print('Done.') + + +def parse_arguments(): + """ + default command line argument parser + :return: args => parsed command line arguments + """ + + parser = argparse.ArgumentParser() + + # parser.add_argument('--config', default='./configs/sample_race_256.yaml') + parser.add_argument("--generator_file", action="store", type=str, + help="pretrained weights file for generator", required=True) + + args = parser.parse_args() + + return args + + +if __name__ == '__main__': + main(parse_arguments()) \ No newline at end of file diff --git a/recognition/45857876/images/20190325144840976.png b/recognition/45857876/images/20190325144840976.png new file mode 100644 index 0000000000..185a3185f2 Binary files /dev/null and b/recognition/45857876/images/20190325144840976.png differ diff --git a/recognition/45857876/images/figure03-style-mixing.png b/recognition/45857876/images/figure03-style-mixing.png new file mode 100644 index 0000000000..a2211c7485 Binary files /dev/null and b/recognition/45857876/images/figure03-style-mixing.png differ diff --git a/recognition/45857876/images/gan.png b/recognition/45857876/images/gan.png new file mode 100644 index 0000000000..76f67d0523 Binary files /dev/null and b/recognition/45857876/images/gan.png differ diff --git a/recognition/45857876/images/logistic_256gen_6_9_1.png b/recognition/45857876/images/logistic_256gen_6_9_1.png new file mode 100644 index 0000000000..b7cf01470e Binary files /dev/null and b/recognition/45857876/images/logistic_256gen_6_9_1.png differ diff --git a/recognition/45857876/images/output512.jpg b/recognition/45857876/images/output512.jpg new file mode 100644 index 0000000000..ccc0d1e97c Binary files /dev/null and b/recognition/45857876/images/output512.jpg differ diff --git a/recognition/45857876/images/rahingegen_6_9_1.png b/recognition/45857876/images/rahingegen_6_9_1.png new file mode 100644 index 0000000000..98afe50603 Binary files /dev/null and b/recognition/45857876/images/rahingegen_6_9_1.png differ diff --git a/recognition/45857876/images/resolution256.gif b/recognition/45857876/images/resolution256.gif new file mode 100644 index 0000000000..6cad4fdf46 Binary files /dev/null and b/recognition/45857876/images/resolution256.gif differ diff --git a/recognition/45857876/images/v2-f1db8c75f4efd04e7eef68b56fefc4d3_1440w.jpg b/recognition/45857876/images/v2-f1db8c75f4efd04e7eef68b56fefc4d3_1440w.jpg new file mode 100644 index 0000000000..44399ec3a9 Binary files /dev/null and b/recognition/45857876/images/v2-f1db8c75f4efd04e7eef68b56fefc4d3_1440w.jpg differ diff --git a/recognition/45857876/model.py b/recognition/45857876/model.py new file mode 100644 index 0000000000..99f4a03989 --- /dev/null +++ b/recognition/45857876/model.py @@ -0,0 +1,762 @@ +import os +import datetime +import time +import timeit +import copy +import random +from collections import OrderedDict + +import torch +import torch.nn as nn +from torch.nn.functional import interpolate + +from utils import get_data_loader +from utils import update_average +from CustomLayers import PixelNormLayer, EqualizedLinear, LayerEpilogue, EqualizedConv2d, BlurLayer, View, StddevLayer + +###### Mapping network +class GMapping(nn.Module): + + def __init__(self, latent_size=512, dlatent_size=512, dlatent_broadcast=None, + mapping_layers=8, mapping_fmaps=512, mapping_lrmul=0.01, mapping_nonlinearity='lrelu', + use_wscale=True, normalize_latents=True, **kwargs): + """ + Mapping network used in the StyleGAN paper. + + """ + + super().__init__() + + self.latent_size = latent_size + self.mapping_fmaps = mapping_fmaps + self.dlatent_size = dlatent_size + self.dlatent_broadcast = dlatent_broadcast + + # Activation function. + act, gain = {'relu': (torch.relu, torch.sqrt(torch.tensor(2))), + 'lrelu': (nn.LeakyReLU(negative_slope=0.2), torch.sqrt(torch.tensor(2)))}[mapping_nonlinearity] + + layers = [] + # Normalize latents. + if normalize_latents: + layers.append(('pixel_norm', PixelNormLayer())) + + # Mapping layers. (apply_bias?) + layers.append(('dense0', EqualizedLinear(self.latent_size, self.mapping_fmaps, + gain=gain, lrmul=mapping_lrmul, use_wscale=use_wscale))) + layers.append(('dense0_act', act)) + for layer_idx in range(1, mapping_layers): + fmaps_in = self.mapping_fmaps + fmaps_out = self.dlatent_size if layer_idx == mapping_layers - 1 else self.mapping_fmaps + layers.append( + ('dense{:d}'.format(layer_idx), + EqualizedLinear(fmaps_in, fmaps_out, gain=gain, lrmul=mapping_lrmul, use_wscale=use_wscale))) + layers.append(('dense{:d}_act'.format(layer_idx), act)) + + # Output. + self.map = nn.Sequential(OrderedDict(layers)) + + def forward(self, x): + # First input: Latent vectors (Z) [mini_batch, latent_size]. + x = self.map(x) + + # Broadcast -> batch_size * dlatent_broadcast * dlatent_size + if self.dlatent_broadcast is not None: + x = x.unsqueeze(1).expand(-1, self.dlatent_broadcast, -1) + return x + + +###### Generator Synthesis block +class GSynthesis(nn.Module): + + def __init__(self, dlatent_size=512, num_channels=3, resolution=1024, + fmap_base=8192, fmap_decay=1.0, fmap_max=512, + use_styles=True, const_input_layer=True, use_noise=True, nonlinearity='lrelu', + use_wscale=True, use_pixel_norm=False, use_instance_norm=True, blur_filter=[1,2,1], + structure='linear', **kwargs): + """ + Synthesis network used in the StyleGAN paper. + + """ + + super().__init__() + + # if blur_filter is None: + # blur_filter = [1, 2, 1] + + def nf(stage): + return min(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_max) + + self.structure = structure + + resolution_log2 = int(torch.log2(torch.tensor(resolution))) + assert resolution == 2 ** resolution_log2 and resolution >= 4 + self.depth = resolution_log2 - 1 + + self.num_layers = resolution_log2 * 2 - 2 + self.num_styles = self.num_layers if use_styles else 1 + + act, gain = {'relu': (torch.relu, torch.sqrt(torch.tensor(2))), + 'lrelu': (nn.LeakyReLU(negative_slope=0.2), torch.sqrt(torch.tensor(2)))}[nonlinearity] + + # Early layers. + self.init_block = InputBlock(nf(1), dlatent_size, const_input_layer, gain, use_wscale, + use_noise, use_pixel_norm, use_instance_norm, use_styles, act) + # create the ToRGB layers for various outputs + rgb_converters = [EqualizedConv2d(nf(1), num_channels, 1, gain=1, use_wscale=use_wscale)] + + # Building blocks for remaining layers. + blocks = [] + for res in range(3, resolution_log2 + 1): + last_channels = nf(res - 2) + channels = nf(res - 1) + # name = '{s}x{s}'.format(s=2 ** res) + blocks.append(GSynthesisBlock(last_channels, channels, blur_filter, dlatent_size, gain, use_wscale, + use_noise, use_pixel_norm, use_instance_norm, use_styles, act)) + rgb_converters.append(EqualizedConv2d(channels, num_channels, 1, gain=1, use_wscale=use_wscale)) + + self.blocks = nn.ModuleList(blocks) + self.to_rgb = nn.ModuleList(rgb_converters) + + # register the temporary upsampler + self.temporaryUpsampler = lambda x: interpolate(x, scale_factor=2) + + def forward(self, dlatents_in, depth=0, alpha=0., labels_in=None): + """ + + """ + + assert depth < self.depth, "Requested output depth cannot be produced" + + if self.structure == 'fixed': + x = self.init_block(dlatents_in[:, 0:2]) + for i, block in enumerate(self.blocks): + x = block(x, dlatents_in[:, 2 * (i + 1):2 * (i + 2)]) + images_out = self.to_rgb[-1](x) + elif self.structure == 'linear': + x = self.init_block(dlatents_in[:, 0:2]) + + if depth > 0: + for i, block in enumerate(self.blocks[:depth - 1]): + x = block(x, dlatents_in[:, 2 * (i + 1):2 * (i + 2)]) + + residual = self.to_rgb[depth - 1](self.temporaryUpsampler(x)) + straight = self.to_rgb[depth](self.blocks[depth - 1](x, dlatents_in[:, 2 * depth:2 * (depth + 1)])) + + images_out = (alpha * straight) + ((1 - alpha) * residual) + else: + images_out = self.to_rgb[0](x) + else: + raise KeyError("Unknown structure: ", self.structure) + + return images_out + +###### Generator block +class Generator(nn.Module): + + def __init__(self, resolution, latent_size=512, dlatent_size=512, + style_mixing_prob=0.9, **kwargs): + """ + # Style-based generator used in the StyleGAN paper. + # Composed of two sub-networks (G_mapping and G_synthesis). + """ + + super(Generator, self).__init__() + + self.style_mixing_prob = style_mixing_prob + + # Setup components. + self.num_layers = (int(torch.log2(torch.tensor(resolution))) - 1) * 2 + self.g_mapping = GMapping(latent_size, dlatent_size, dlatent_broadcast=self.num_layers, **kwargs) + self.g_synthesis = GSynthesis(resolution=resolution, **kwargs) + + def forward(self, latents_in, depth, alpha): + """ + + """ + + dlatents_in = self.g_mapping(latents_in) + + if self.training: + # Update moving average of W(dlatent). + # if self.truncation is not None: + # self.truncation.update(dlatents_in[0, 0].detach()) + + # Perform style mixing regularization. + if self.style_mixing_prob is not None and self.style_mixing_prob > 0: + latents2 = torch.randn(latents_in.shape).to(latents_in.device) + dlatents2 = self.g_mapping(latents2) + layer_idx = torch.unsqueeze(torch.arange(self.num_layers), 1).to( + latents_in.device) + cur_layers = 2 * (depth + 1) + mixing_cutoff = random.randint(1, + cur_layers) if random.random() < self.style_mixing_prob else cur_layers + dlatents_in = torch.where(layer_idx < mixing_cutoff, dlatents_in, dlatents2) + + fake_images = self.g_synthesis(dlatents_in, depth, alpha) + + return fake_images + + + +###### Discriminator block +class Discriminator(nn.Module): + + def __init__(self, resolution, num_channels=3, fmap_base=8192, fmap_decay=1.0, fmap_max=512, + nonlinearity='lrelu', use_wscale=True, mbstd_group_size=4, mbstd_num_features=1, + blur_filter=None, structure='linear', **kwargs): + """ + Discriminator used in the StyleGAN paper. + + """ + super(Discriminator, self).__init__() + + def nf(stage): + return min(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_max) + + self.mbstd_num_features = mbstd_num_features + self.mbstd_group_size = mbstd_group_size + self.structure = structure + + resolution_log2 = int(torch.log2(torch.tensor(resolution))) + assert resolution == 2 ** resolution_log2 and resolution >= 4 + self.depth = resolution_log2 - 1 + + act, gain = {'relu': (torch.relu, torch.sqrt(torch.tensor(2))), + 'lrelu': (nn.LeakyReLU(negative_slope=0.2), torch.sqrt(torch.tensor(2)))}[nonlinearity] + + # create the remaining layers + blocks = [] + from_rgb = [] + for res in range(resolution_log2, 2, -1): + blocks.append(DiscriminatorBlock(nf(res - 1), nf(res - 2), + gain=gain, use_wscale=use_wscale, activation_layer=act, + blur_kernel=blur_filter)) + # create the fromRGB layers for various inputs: + from_rgb.append(EqualizedConv2d(num_channels, nf(res - 1), kernel_size=1, + gain=gain, use_wscale=use_wscale)) + self.blocks = nn.ModuleList(blocks) + + # Building the final block. + self.final_block = DiscriminatorTop(self.mbstd_group_size, self.mbstd_num_features, + in_channels=nf(2), intermediate_channels=nf(2), + gain=gain, use_wscale=use_wscale, activation_layer=act) + from_rgb.append(EqualizedConv2d(num_channels, nf(2), kernel_size=1, + gain=gain, use_wscale=use_wscale)) + self.from_rgb = nn.ModuleList(from_rgb) + + # register the temporary downSampler + self.temporaryDownsampler = nn.AvgPool2d(2) + + def forward(self, images_in, depth, alpha=1., labels_in=None): + """ + """ + + assert depth < self.depth, "Requested output depth cannot be produced" + + if self.structure == 'fixed': + x = self.from_rgb[0](images_in) + for i, block in enumerate(self.blocks): + x = block(x) + scores_out = self.final_block(x) + elif self.structure == 'linear': + if depth > 0: + residual = self.from_rgb[self.depth - depth](self.temporaryDownsampler(images_in)) + straight = self.blocks[self.depth - depth - 1](self.from_rgb[self.depth - depth - 1](images_in)) + x = (alpha * straight) + ((1 - alpha) * residual) + + for block in self.blocks[(self.depth - depth):]: + x = block(x) + else: + x = self.from_rgb[-1](images_in) + + scores_out = self.final_block(x) + else: + raise KeyError("Unknown structure: ", self.structure) + + return scores_out + +###### Constant Input block +class InputBlock(nn.Module): + """ + The fixed constant as input for 4*4 resolution + """ + + def __init__(self, nf, dlatent_size, const_input_layer, gain, + use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer): + super().__init__() + self.const_input_layer = const_input_layer + self.nf = nf + + if self.const_input_layer: + # called 'const' in torch + self.const = nn.Parameter(torch.ones(1, nf, 4, 4)) + self.bias = nn.Parameter(torch.ones(nf)) + else: + self.dense = EqualizedLinear(dlatent_size, nf * 16, gain=gain / 4, + use_wscale=use_wscale) + # tweak gain to match the official implementation of Progressing GAN + + self.epi1 = LayerEpilogue(nf, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, + use_styles, activation_layer) + self.conv = EqualizedConv2d(nf, nf, 3, gain=gain, use_wscale=use_wscale) + self.epi2 = LayerEpilogue(nf, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, + use_styles, activation_layer) + + def forward(self, dlatents_in_range): + batch_size = dlatents_in_range.size(0) + + if self.const_input_layer: + x = self.const.expand(batch_size, -1, -1, -1) + x = x + self.bias.view(1, -1, 1, 1) + else: + x = self.dense(dlatents_in_range[:, 0]).view(batch_size, self.nf, 4, 4) + + x = self.epi1(x, dlatents_in_range[:, 0]) + x = self.conv(x) + x = self.epi2(x, dlatents_in_range[:, 1]) + + return x + +###### Generator Synthesis Block +class GSynthesisBlock(nn.Module): + def __init__(self, in_channels, out_channels, blur_filter, dlatent_size, gain, + use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer): + # 2**res x 2**res + # res = 3..resolution_log2 + super().__init__() + + if blur_filter: + blur = BlurLayer(blur_filter) + else: + blur = None + + self.conv0_up = EqualizedConv2d(in_channels, out_channels, kernel_size=3, gain=gain, use_wscale=use_wscale, + intermediate=blur, upscale=True) + self.epi1 = LayerEpilogue(out_channels, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, + use_styles, activation_layer) + self.conv1 = EqualizedConv2d(out_channels, out_channels, kernel_size=3, gain=gain, use_wscale=use_wscale) + self.epi2 = LayerEpilogue(out_channels, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, + use_styles, activation_layer) + + def forward(self, x, dlatents_in_range): + x = self.conv0_up(x) + x = self.epi1(x, dlatents_in_range[:, 0]) + x = self.conv1(x) + x = self.epi2(x, dlatents_in_range[:, 1]) + return x + +###### Discriminator top Block +class DiscriminatorTop(nn.Sequential): + def __init__(self, + mbstd_group_size, + mbstd_num_features, + in_channels, + intermediate_channels, + gain, use_wscale, + activation_layer, + resolution=4, + in_channels2=None, + output_features=1, + last_gain=1): + """ + + """ + + layers = [] + if mbstd_group_size > 1: + layers.append(('stddev_layer', StddevLayer(mbstd_group_size, mbstd_num_features))) + + if in_channels2 is None: + in_channels2 = in_channels + + layers.append(('conv', EqualizedConv2d(in_channels + mbstd_num_features, in_channels2, kernel_size=3, + gain=gain, use_wscale=use_wscale))) + layers.append(('act0', activation_layer)) + layers.append(('view', View(-1))) + layers.append(('dense0', EqualizedLinear(in_channels2 * resolution * resolution, intermediate_channels, + gain=gain, use_wscale=use_wscale))) + layers.append(('act1', activation_layer)) + layers.append(('dense1', EqualizedLinear(intermediate_channels, output_features, + gain=last_gain, use_wscale=use_wscale))) + + super().__init__(OrderedDict(layers)) + + +###### Discriminator Block(wrap discrimnator together) +class DiscriminatorBlock(nn.Sequential): + def __init__(self, in_channels, out_channels, gain, use_wscale, activation_layer, blur_kernel): + super().__init__(OrderedDict([ + ('conv0', EqualizedConv2d(in_channels, in_channels, kernel_size=3, gain=gain, use_wscale=use_wscale)), + # out channels nf(res-1) + ('act0', activation_layer), + ('blur', BlurLayer(kernel=blur_kernel)), + ('conv1_down', EqualizedConv2d(in_channels, out_channels, kernel_size=3, + gain=gain, use_wscale=use_wscale, downscale=True)), + ('act1', activation_layer)])) + + +###### Loss function of logistic +class LogisticGAN: + def __init__(self, dis): + self.dis = dis + + # gradient penalty + def R1Penalty(self, real_img, height, alpha): + + apply_loss_scaling = lambda x: x * torch.exp(x * [torch.log(torch.tensor(2.0))].to(real_img.device)) + undo_loss_scaling = lambda x: x * torch.exp(-x * [torch.log(torch.tensor(2.0))].to(real_img.device)) + + real_img = torch.autograd.Variable(real_img, requires_grad=True) + real_logit = self.dis(real_img, height, alpha) + + real_grads = torch.autograd.grad(outputs=real_logit, inputs=real_img, + grad_outputs=torch.ones(real_logit.size()).to(real_img.device), + create_graph=True, retain_graph=True)[0].view(real_img.size(0), -1) + + r1_penalty = torch.sum(torch.mul(real_grads, real_grads)) + return r1_penalty + + def dis_loss(self, real_samps, fake_samps, height, alpha, r1_gamma=10.0): + # Obtain predictions + r_preds = self.dis(real_samps, height, alpha) + f_preds = self.dis(fake_samps, height, alpha) + + loss = torch.mean(nn.Softplus()(f_preds)) + torch.mean(nn.Softplus()(-r_preds)) + + if r1_gamma != 0.0: + r1_penalty = self.R1Penalty(real_samps.detach(), height, alpha) * (r1_gamma * 0.5) + loss += r1_penalty + + return loss + + def gen_loss(self, real_samps, fake_samps, height, alpha): + f_preds = self.dis(fake_samps, height, alpha) + + return torch.mean(nn.Softplus()(-f_preds)) + + +###### Loss function of Relativistic Average Hinge +class RelativisticAverageHingeGAN: + + def __init__(self, dis): + self.dis = dis + + def dis_loss(self, real_samps, fake_samps, height, alpha): + # Obtain predictions + r_preds = self.dis(real_samps, height, alpha) + f_preds = self.dis(fake_samps, height, alpha) + + # difference between real and fake: + r_f_diff = r_preds - torch.mean(f_preds) + + # difference between fake and real samples + f_r_diff = f_preds - torch.mean(r_preds) + + # return the loss + loss = (torch.mean(nn.ReLU()(1 - r_f_diff)) + + torch.mean(nn.ReLU()(1 + f_r_diff))) + + return loss + + def gen_loss(self, real_samps, fake_samps, height, alpha): + # Obtain predictions + r_preds = self.dis(real_samps, height, alpha) + f_preds = self.dis(fake_samps, height, alpha) + + # difference between real and fake: + r_f_diff = r_preds - torch.mean(f_preds) + + # difference between fake and real samples + f_r_diff = f_preds - torch.mean(r_preds) + + # return the loss + return (torch.mean(nn.ReLU()(1 + r_f_diff)) + + torch.mean(nn.ReLU()(1 - f_r_diff))) + + +###### stylegan class (wrap generator and diacriminator together) +class StyleGAN: + + def __init__(self, structure, resolution, num_channels, latent_size,loss,drift=0.001, + d_repeats=1, use_ema=False, ema_decay=0.999, device=torch.device("cpu")): + """ + Wrapper around the Generator and the Discriminator. + + """ + + # state of the object + assert structure in ['fixed', 'linear'] + self.structure = structure + self.depth = int(torch.log2(torch.tensor(resolution))) - 1 + self.latent_size = latent_size + self.device = device + self.d_repeats = d_repeats + + self.use_ema = use_ema + self.ema_decay = ema_decay + + # Create the Generator and the Discriminator + + self.gen = Generator(num_channels=3, + resolution=resolution, + structure=self.structure + ).to(self.device) + + self.dis = Discriminator(num_channels=3, + resolution=resolution, + structure=self.structure + ).to(self.device) + + # define the optimizers for the discriminator and generator + self.__setup_gen_optim() + self.__setup_dis_optim() + + # define the loss function used for training the GAN + self.drift = drift + self.loss = self.__setup_loss(loss) + + # Use of ema + if self.use_ema: + # create a shadow copy of the generator + self.gen_shadow = copy.deepcopy(self.gen) + # updater function: + self.ema_updater = update_average + # initialize the gen_shadow weights equal to the weights of gen + self.ema_updater(self.gen_shadow, self.gen, beta=0) + + def __setup_gen_optim(self): + self.gen_optim = torch.optim.Adam(self.gen.parameters(), lr=0.003, betas=(0, 0.99), eps=1e-8) + + def __setup_dis_optim(self): + self.dis_optim = torch.optim.Adam(self.dis.parameters(),lr=0.003, betas=(0, 0.99), eps=1e-8) + + def __setup_loss(self,loss): + if loss == "logistic": + loss = LogisticGAN(self.dis) + elif loss == "RAhinge": + loss = RelativisticAverageHingeGAN(self.dis) + + return loss + + def __progressive_down_sampling(self, real_batch, depth, alpha): + """ + private helper for down_sampling the original images in order to facilitate the + progressive growing of the layers. + + """ + + from torch.nn import AvgPool2d + from torch.nn.functional import interpolate + + if self.structure == 'fixed': + return real_batch + + # down_sample the real_batch for the given depth + down_sample_factor = int(torch.pow(torch.tensor(2), torch.tensor(self.depth - depth - 1))) + prior_down_sample_factor = max(int(torch.pow(torch.tensor(2), torch.tensor(self.depth - depth))), 0) + + ds_real_samples = AvgPool2d(down_sample_factor)(real_batch) + + if depth > 0: + prior_ds_real_samples = interpolate(AvgPool2d(prior_down_sample_factor)(real_batch), scale_factor=2) + else: + prior_ds_real_samples = ds_real_samples + + # real samples are a combination of ds_real_samples and prior_ds_real_samples + real_samples = (alpha * ds_real_samples) + ((1 - alpha) * prior_ds_real_samples) + + # return the so computed real_samples + return real_samples + + def optimize_discriminator(self, noise, real_batch, depth, alpha): + """ + performs one step of weight update on discriminator using the batch of data + """ + + real_samples = self.__progressive_down_sampling(real_batch, depth, alpha) + + loss_val = 0 + for _ in range(self.d_repeats): + # generate a batch of samples + fake_samples = self.gen(noise, depth, alpha).detach() + + loss = self.loss.dis_loss(real_samples, fake_samples, depth, alpha) + + # optimize discriminator + self.dis_optim.zero_grad() + loss.backward() + self.dis_optim.step() + + loss_val += loss.item() + + return loss_val / self.d_repeats + + def optimize_generator(self, noise, real_batch, depth, alpha): + """ + performs one step of weight update on generator for the given batch_size + + """ + + real_samples = self.__progressive_down_sampling(real_batch, depth, alpha) + + # generate fake samples: + fake_samples = self.gen(noise, depth, alpha) + + # Change this implementation for making it compatible for relativisticGAN + loss = self.loss.gen_loss(real_samples, fake_samples, depth, alpha) + + # optimize the generator + self.gen_optim.zero_grad() + loss.backward() + # Gradient Clipping + nn.utils.clip_grad_norm_(self.gen.parameters(), max_norm=10.) + self.gen_optim.step() + + # if use_ema is true, apply ema to the generator parameters + if self.use_ema: + self.ema_updater(self.gen_shadow, self.gen, self.ema_decay) + + # return the loss value + return loss.item() + + @staticmethod + def create_grid(samples, scale_factor, img_file): + """ + utility function to create a grid of GAN samples + + """ + from torchvision.utils import save_image + from torch.nn.functional import interpolate + + # upsample the image + if scale_factor > 1: + samples = interpolate(samples, scale_factor=scale_factor) + + # save the images: + save_image(samples, img_file, nrow=int(torch.sqrt(torch.tensor(len(samples)))), + normalize=True, scale_each=True, pad_value=128, padding=1) + + def train(self, dataset, num_workers, epochs, batch_sizes, fade_in_percentage, logger, output, + num_samples=36, start_depth=0, feedback_factor=100, checkpoint_factor=1): + """ + Utility method for training the GAN. + + """ + + assert self.depth <= len(epochs), "epochs not compatible with depth" + assert self.depth <= len(batch_sizes), "batch_sizes not compatible with depth" + assert self.depth <= len(fade_in_percentage), "fade_in_percentage not compatible with depth" + + # turn the generator and discriminator into train mode + self.gen.train() + self.dis.train() + if self.use_ema: + self.gen_shadow.train() + + # create a global time counter + global_time = time.time() + + # create fixed_input for debugging + fixed_input = torch.randn(num_samples, self.latent_size).to(self.device) + + # config depend on structure + logger.info("Starting the training process ... \n") + if self.structure == 'fixed': + start_depth = self.depth - 1 + step = 1 # counter for number of iterations + for current_depth in range(start_depth, self.depth): + current_res = torch.pow(torch.tensor(2), torch.tensor(current_depth + 2)) + logger.info("Currently working on depth: %d", current_depth + 1) + logger.info("Current resolution: %d x %d" % (int(current_res), int(current_res))) + + ticker = 1 + + # Choose training parameters and configure training ops. + # TODO + data = get_data_loader(dataset, batch_sizes[current_depth], num_workers) + + for epoch in range(1, epochs[current_depth] + 1): + start = timeit.default_timer() # record time at the start of epoch + + logger.info("Epoch: [%d]" % epoch) + # total_batches = len(iter(data)) + total_batches = len(data) + + fade_point = int((fade_in_percentage[current_depth] / 100) + * epochs[current_depth] * total_batches) + + for (i, batch) in enumerate(data, 1): + # calculate the alpha for fading in the layers + alpha = ticker / fade_point if ticker <= fade_point else 1 + + # extract current batch of data for training + images = batch.to(self.device) + gan_input = torch.randn(images.shape[0], self.latent_size).to(self.device) + + # optimize the discriminator: + dis_loss = self.optimize_discriminator(gan_input, images, current_depth, alpha) + + # optimize the generator: + gen_loss = self.optimize_generator(gan_input, images, current_depth, alpha) + + # provide a loss feedback + if i % int(total_batches / feedback_factor + 1) == 0 or i == 1: + elapsed = time.time() - global_time + elapsed = str(datetime.timedelta(seconds=elapsed)).split('.')[0] + logger.info( + "Elapsed: [%s] Step: %d Batch: %d D_Loss: %f G_Loss: %f" + % (elapsed, step, i, dis_loss, gen_loss)) + + # create a grid of samples and save it + os.makedirs(os.path.join(output, 'samples'), exist_ok=True) + gen_img_file = os.path.join(output, 'samples', "gen_" + str(current_depth) + + "_" + str(epoch) + "_" + str(i) + ".png") + + with torch.no_grad(): + self.create_grid( + samples=self.gen(fixed_input, current_depth, alpha).detach() if not self.use_ema + else self.gen_shadow(fixed_input, current_depth, alpha).detach(), + scale_factor=int( + torch.pow(torch.tensor(2), torch.tensor( + self.depth - current_depth - 1))) if self.structure == 'linear' else 1, + img_file=gen_img_file, + ) + + # increment the alpha ticker and the step + ticker += 1 + step += 1 + + elapsed = timeit.default_timer() - start + elapsed = str(datetime.timedelta(seconds=elapsed)).split('.')[0] + logger.info("Time taken for epoch: %s\n" % elapsed) + logger.info("Current Alpha is %f" % alpha) + + if epoch % checkpoint_factor == 0 or epoch == 1 or epoch == epochs[current_depth]: + save_dir = os.path.join(output, 'models') + os.makedirs(save_dir, exist_ok=True) + gen_save_file = os.path.join(save_dir, "GAN_GEN_" + str(current_depth) + "_" + str(epoch) + ".pth") + dis_save_file = os.path.join(save_dir, "GAN_DIS_" + str(current_depth) + "_" + str(epoch) + ".pth") + gen_optim_save_file = os.path.join( + save_dir, "GAN_GEN_OPTIM_" + str(current_depth) + "_" + str(epoch) + ".pth") + dis_optim_save_file = os.path.join( + save_dir, "GAN_DIS_OPTIM_" + str(current_depth) + "_" + str(epoch) + ".pth") + + torch.save(self.gen.state_dict(), gen_save_file) + logger.info("Saving the model to: %s\n" % gen_save_file) + torch.save(self.dis.state_dict(), dis_save_file) + torch.save(self.gen_optim.state_dict(), gen_optim_save_file) + torch.save(self.dis_optim.state_dict(), dis_optim_save_file) + + # also save the shadow generator if use_ema is True + if self.use_ema: + gen_shadow_save_file = os.path.join( + save_dir, "GAN_GEN_SHADOW_" + str(current_depth) + "_" + str(epoch) + ".pth") + torch.save(self.gen_shadow.state_dict(), gen_shadow_save_file) + logger.info("Saving the model to: %s\n" % gen_shadow_save_file) + + logger.info('Training completed.\n') + + +if __name__ == '__main__': + print('Done.') \ No newline at end of file diff --git a/recognition/45857876/readme.md b/recognition/45857876/readme.md new file mode 100644 index 0000000000..30d20f0e9d --- /dev/null +++ b/recognition/45857876/readme.md @@ -0,0 +1,140 @@ +# **Knee MRI Image stylegan** +Student name: Mengyao Ma +Student number: s4585787 + +### Generative Adversarial Network +![Gan structure](https://github.com/MMMMMYY/PatternFlow/blob/topic-recognition/recognition/45857876/images/gan.png) +**GAN** is short for **Generative Adversarial Network** proposed by Ian Goodfellow in 2014. The main structure of GAN includes a **Generator(G)** and a **Discriminator (D)**. For the generator, the input requires an n-dimensional vector, and the output is a picture with the pixel size of the picture, while the discriminator discriminates the pictures generated by the generator and labels them as "fake" or "real". + + +### StyleGAN +**StyleGAN** is inspired by **style transfer** to design a **new generator structure**. In addtion, StyleGAN is **evolved** from **ProGAN** and uses a similar network structure. StyleGAN uses style to affect the posture and identity characteristics of the face, and noise to affect details such as hair strands, wrinkles, and skin tone. + + +## The main StyleGAN structure +![Traditional genertaor and StyleGAN generator](https://github.com/MMMMMYY/PatternFlow/blob/topic-recognition/recognition/45857876/images/v2-f1db8c75f4efd04e7eef68b56fefc4d3_1440w.jpg) +Through observation, it can be found that different layers and resolutions will affect different features. The lower the layer and resolution, the coarser the features it affects. We can divide these characteristics into three types: +1. Rough-(resolution 0-8), affecting posture, general hairstyle, facial shape, etc.; +2. Medium-(resolution 16-32), affecting finer facial features, hairstyles, opening of eyes or Closed, etc.; +3. High-quality-(resolution 64-1024), affecting color (eyes, hair and skin) and microscopic features; + + +In traditional gan, the generator only feeds the random variables as latent code into input layer. +### Mapping Network + +To better solve the Feature unwrapping problem, the Mapping network is added into Generator. It consists of 8 fully connected layers. The Mapping network convert the latent code into 18 vectors after affin transform. These vectors A learned affin transform are added into the synthesis generator network(two for each resolution). The vectors A can control the style on certain resolution. + +### AdaIN module + +![AdaIN](https://github.com/MMMMMYY/PatternFlow/blob/topic-recognition/recognition/45857876/images/20190325144840976.png) + +At each resolution, two A will affect the generator twice, once after Upsampling and once after Convolution by using AdaIN. +Like the equation ablow, expand A into scaling factors y~𝑠,𝑖~ and deviation factors y~𝑏,𝑖~, and make a weighted sum of these two factors and the normalized convolution output to apply an influence. + +### Random nosie + +To control the character detail and diversity of image generated, add a scaled noise to each channel before the Adain module. + +### Style mixing + +To decrease the correlation, the genarator will randomly choose two vector A, and mix them by mixing factor(in defualt setting, the factor is 0.5). The randomly switch ensure the network will not reply on either correlation between each blocks. + +## **Loss function applied** + +In this task, I explore two loss functions to find the influence of loss function. + +### 1. Logistic loss function +Output for logistic loss function for resolution 256 shows below: +![resolution 256 epoch9](https://github.com/MMMMMYY/PatternFlow/blob/topic-recognition/recognition/45857876/images/logistic_256gen_6_9_1.png) + +### 2. Relativistic Average Hinge loss function +Output for logistic loss function for resolution 256 shows below: +![resolution 256 epoch9](https://github.com/MMMMMYY/PatternFlow/blob/topic-recognition/recognition/45857876/images/rahingegen_6_9_1.png) + + +The two models setting is same except loss function. The result image above is the result of epoch 6 the resolution is 256. Because it is too time consuming, so I just pick the result of epoch 6. +**However, it is easy to find when the Relativistic Average Hinge loss function is applied, the deformation at the edge of the knee image is more serious. Although, the image generated with Logistic loss fuction also has some strange pixels, it performs better than image with Relativistic Average Hinge loss function.** + + + +## **Output of each resolution** +**The result will present images with Logistic loss function.** +![resolution 4-512](https://github.com/MMMMMYY/PatternFlow/blob/topic-recognition/recognition/45857876/images/output512.jpg) +![resolution for all precedure](https://github.com/MMMMMYY/PatternFlow/blob/topic-recognition/recognition/45857876/images/resolution256.gif) + +From the result we can find, the details will change randomly as the number of layers increases cause the random nosie added in every layer. Compared with the result of resolution 256, knee edges and details of the preview result 512 are much clearer, and there are no obvious error pixels. + +## **Style mixing output** + +![test_image](https://github.com/MMMMMYY/PatternFlow/blob/topic-recognition/recognition/45857876/images/figure03-style-mixing.png) + + +The test image is generated by generator model (6_32) which is resolution 256 and epoch 32. Because I did not finish training resolution 512, the generator preforms bad on resolution 512. +## Requirments + +yacs +tqdm +numpy (only for visualization, file test_script.py, generate_mixing_figure.py) +torchvision +torch + +## Execute the code +Training: + python train.py +Generate mixing image(test): + python generate_mixing_figure.py --generator_file [new256/models/GAN_GEN_4_16.pth](for example) + + +## The default setting: + +general setting: + device = 'cuda' + number of preview samples = 36 + make checkpoint every/epoch = 10 + test loss: + epochs: [2,4,8,8,16,24,32] + style mixing: + epochs: [2,4,8,8,16,24,32,40] + +Generator setting: + latent size = 512 + mapping layers = 4(8 in original paper, but 4 layers when latene size = 512) + blur_filter = [1, 2, 1] + +Discriminator setting: + enable equalized learning rate = True + blur_filter = [1, 2, 1] + +Generator Optimizer setting: + optimizer = Adam + learning_rate = 0.003 + betas = [0,0.99] + eps = defualt + +Discriminator Optimizer setting: + optimizer = Adam + learning_rate = 0.003 + betas = [0,0.99] + eps = defualt + + + +## Reference +Paper: +A Style-Based Generator Architecture for Generative Adversarial Networks +Tero Karras (NVIDIA), Samuli Laine (NVIDIA), Timo Aila (NVIDIA) +https://arxiv.org/abs/1812.04948 +Progressive Growing of GANs for Improved Quality, Stability, and Variation +Tero Karras, Timo Aila, Samuli Laine, Jaakko Lehtinen +https://arxiv.org/abs/1710.10196 +Generative Adversarial Networks +Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, Yoshua Bengio +https://arxiv.org/abs/1406.2661 + +Code: +https://github.com/huangzh13/StyleGAN.pytorch +https://github.com/lernapparat/lernapparat +https://github.com/NVlabs/stylegan +https://github.com/akanimax/pro_gan_pytorch +https://github.com/rosinality/style-based-gan-pytorch +https://github.com/goodfeli/adversarial diff --git a/recognition/45857876/train.py b/recognition/45857876/train.py new file mode 100644 index 0000000000..7d131c1aa1 --- /dev/null +++ b/recognition/45857876/train.py @@ -0,0 +1,127 @@ +import os +import argparse +import shutil + +import torch +from torch.backends import cudnn + +from utils import make_dataset,make_logger, list_dir_recursively_with_ignore, copy_files_and_create_dirs +from model import StyleGAN + +output_dir = '/stylegan/rahinge256' +device = "cuda" +device_id = "0" +resolution = 256 +use_ema = True + + +# Load fewer layers of pre-trained models if possible +def load(model, cpk_file): + pretrained_dict = torch.load(cpk_file) + model_dict = model.state_dict() + pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict} + model_dict.update(pretrained_dict) + model.load_state_dict(model_dict) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description="StyleGAN pytorch implementation.") + parser.add_argument('--config', default='./configs/sample.yaml') + + parser.add_argument("--start_depth", action="store", type=int, default=0, + help="Starting depth for training the network") + + parser.add_argument("--generator_file", action="store", type=str, default=None, + help="pretrained Generator file (compatible with my code)") + parser.add_argument("--gen_shadow_file", action="store", type=str, default=None, + help="pretrained gen_shadow file") + parser.add_argument("--discriminator_file", action="store", type=str, default=None, + help="pretrained Discriminator file (compatible with my code)") + parser.add_argument("--gen_optim_file", action="store", type=str, default=None, + help="saved state of generator optimizer") + parser.add_argument("--dis_optim_file", action="store", type=str, default=None, + help="saved_state of discriminator optimizer") + args = parser.parse_args() + + + # make output dir + if os.path.exists(output_dir): + raise KeyError("Existing path: ", output_dir) + os.makedirs(output_dir) + + print("copy") + # copy codes and config file + files = list_dir_recursively_with_ignore('.', ignores=['diagrams', 'configs']) + + files = [(f[0], os.path.join(output_dir, "src", f[1])) for f in files] + + # copy_files_and_create_dirs(files) + + # shutil.copy2(args.config, output_dir) + print("finish copy") + # logger + logger = make_logger("project", output_dir, 'log') + + # device + if device == 'cuda': + os.environ['CUDA_VISIBLE_DEVICES'] = device_id + num_gpus = len(device_id.split(',')) + logger.info("Using {} GPUs.".format(num_gpus)) + logger.info("Training on {}.\n".format(torch.cuda.get_device_name(0))) + cudnn.benchmark = True + device = torch.device(device) + + # create the dataset for training + dataset = make_dataset("AKOA_Analysis",256) + + # init the network + style_gan = StyleGAN(structure="linear", + resolution= resolution, + num_channels= 3, + latent_size= 512, + loss = "RAhinge", + drift=0.001, + d_repeats=1, + use_ema=True, + ema_decay=0.999, + device=device) + + # Resume training from checkpoints + if args.generator_file is not None: + logger.info("Loading generator from: %s", args.generator_file) + # style_gan.gen.load_state_dict(torch.load(args.generator_file)) + # Load fewer layers of pre-trained models if possible + load(style_gan.gen, args.generator_file) + else: + logger.info("Training from scratch...") + + if args.discriminator_file is not None: + logger.info("Loading discriminator from: %s", args.discriminator_file) + style_gan.dis.load_state_dict(torch.load(args.discriminator_file)) + + if args.gen_shadow_file is not None and use_ema: + logger.info("Loading shadow generator from: %s", args.gen_shadow_file) + # style_gan.gen_shadow.load_state_dict(torch.load(args.gen_shadow_file)) + # Load fewer layers of pre-trained models if possible + load(style_gan.gen_shadow, args.gen_shadow_file) + + if args.gen_optim_file is not None: + logger.info("Loading generator optimizer from: %s", args.gen_optim_file) + style_gan.gen_optim.load_state_dict(torch.load(args.gen_optim_file)) + + if args.dis_optim_file is not None: + logger.info("Loading discriminator optimizer from: %s", args.dis_optim_file) + style_gan.dis_optim.load_state_dict(torch.load(args.dis_optim_file)) + + # train the network + style_gan.train(dataset=dataset, + num_workers=4, + epochs=[2,4,8,8,16,24,32], + batch_sizes=[128, 128, 128, 64, 32, 16, 8], + fade_in_percentage=[50, 50, 50, 50, 50, 50, 50], + logger=logger, + output=output_dir, + num_samples=36, + start_depth= 0, + feedback_factor=10, + checkpoint_factor=10) diff --git a/recognition/45857876/utils.py b/recognition/45857876/utils.py new file mode 100644 index 0000000000..9cf9d5568b --- /dev/null +++ b/recognition/45857876/utils.py @@ -0,0 +1,201 @@ +from torchvision.transforms import ToTensor, Normalize, Compose, Resize, RandomHorizontalFlip +from torch.utils.data import Dataset +from torch.utils.data import DataLoader +from PIL import Image +import fnmatch +import os +import shutil +from typing import List, Tuple +import logging +import sys + + +####### transform image +def get_transform(new_size=None): + """ + obtain the image transforms required for the input data + """ + + if new_size is not None: + image_transform = Compose([ + RandomHorizontalFlip(), + Resize(new_size), + ToTensor(), + Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) + ]) + + else: + image_transform = Compose([ + RandomHorizontalFlip(), + ToTensor(), + Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) + ]) + return image_transform + + +###### Make dataset +def make_dataset(data_dir, resolution): + """ + generate the dataset + """ + _dataset = FlatDirectoryImageDataset(data_dir, transform=get_transform(new_size=(resolution, resolution))) + return _dataset + +###### Data loader +def get_data_loader(dataset, batch_size, num_workers): + """ + generate the data_loader from the given dataset + """ + + dl = DataLoader( + dataset, + batch_size=batch_size, + shuffle=True, + num_workers=num_workers, + drop_last=True, + pin_memory=True + ) + + return dl + +###### Wrap data dataset +class FlatDirectoryImageDataset(Dataset): + """ pyTorch Dataset wrapper for the generic flat directory images dataset """ + + def __setup_files(self): + """ + private helper for setting up the files_list + """ + file_names = os.listdir(self.data_dir) + files = [] # initialize to empty list + + for file_name in file_names: + possible_file = os.path.join(self.data_dir, file_name) + if os.path.isfile(possible_file): + files.append(possible_file) + + # return the files list + return files + + def __init__(self, data_dir, transform=None): + """ + constructor for the class + """ + # define the state of the object + self.data_dir = data_dir + self.transform = transform + + # setup the files for reading + self.files = self.__setup_files() + + def __len__(self): + """ + compute the length of the dataset + """ + return len(self.files) + + def __getitem__(self, idx): + """ + obtain the image (read and transform) + """ + img = Image.open(self.files[idx]).convert('RGB') + img = self.transform(img) + return img + +######update the model target +def update_average(model_tgt, model_src, beta): + """ + update the model_target using exponential moving averages + """ + + # utility function for toggling the gradient requirements of the models + def toggle_grad(model, requires_grad): + for p in model.parameters(): + p.requires_grad_(requires_grad) + + # turn off gradient calculation + toggle_grad(model_tgt, False) + toggle_grad(model_src, False) + + param_dict_src = dict(model_src.named_parameters()) + + for p_name, p_tgt in model_tgt.named_parameters(): + p_src = param_dict_src[p_name] + assert (p_src is not p_tgt) + p_tgt.copy_(beta * p_tgt + (1. - beta) * p_src) + + # turn back on the gradient calculation + toggle_grad(model_tgt, True) + toggle_grad(model_src, True) + + + +#####copy file for backup +def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None: + """Takes in a list of tuples of (src, dst) paths and copies files. + Will create all necessary directories.""" + for file in files: + target_dir_name = os.path.dirname(file[1]) + + # will create all intermediate-level directories + if not os.path.exists(target_dir_name): + os.makedirs(target_dir_name) + + shutil.copyfile(file[0], file[1]) + +######list the file path +def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> \ + List[Tuple[str, str]]: + """List all files recursively in a given directory while ignoring given file and directory names. + Returns list of tuples containing both absolute and relative paths.""" + assert os.path.isdir(dir_path) + base_name = os.path.basename(os.path.normpath(dir_path)) + + if ignores is None: + ignores = [] + + result = [] + + for root, dirs, files in os.walk(dir_path, topdown=True): + for ignore_ in ignores: + dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)] + + # dirs need to be edited in-place + for d in dirs_to_remove: + dirs.remove(d) + + files = [f for f in files if not fnmatch.fnmatch(f, ignore_)] + + absolute_paths = [os.path.join(root, f) for f in files] + relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths] + + if add_base_to_relative: + relative_paths = [os.path.join(base_name, p) for p in relative_paths] + + assert len(absolute_paths) == len(relative_paths) + result += zip(absolute_paths, relative_paths) + + return result + + +###### make logger of whole training and save +def make_logger(name, save_dir, save_filename): + DATE_FORMAT = "%Y-%m-%d %H:%M:%S" + + logger = logging.getLogger(name) + logger.setLevel(logging.DEBUG) + + ch = logging.StreamHandler(stream=sys.stdout) + ch.setLevel(logging.DEBUG) + formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s", datefmt=DATE_FORMAT) + ch.setFormatter(formatter) + + logger.addHandler(ch) + + if save_dir: + fh = logging.FileHandler(os.path.join(save_dir, save_filename + ".txt"), mode='w') + fh.setLevel(logging.DEBUG) + fh.setFormatter(formatter) + logger.addHandler(fh) + + return logger \ No newline at end of file