{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Introduction\n### A few words\n\nSubmitted versions are versions 4 and 5 that score **30.7 and 30.3 on public**. The only difference is that I used different class sampling probabilities for generating images. The solution proposed here is legit, no memorization of any kind is made.\n\nIf you wish to evaluate the results locally, make sure to add the metric dataset, change the PATH varaibles and uncomment the scoring lines.\n\n### References\n\n- NVIDIA's Progressive Growing of GANs paper : https://research.nvidia.com/publication/2017-10_Progressive-Growing-of\n- CGANs with Projection Discriminator : https://arxiv.org/pdf/1802.05637.pdf\n- The modeling part of the kernel is taken from : https://github.com/akanimax/pro_gan_pytorch"},{"metadata":{},"cell_type":"markdown","source":"# Imports"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import re\nimport gc\nimport os\nimport cv2\nimport copy\nimport time\nimport pickle\nimport random\nimport shutil\nimport urllib\nimport pathlib\nimport datetime\nimport operator\nimport warnings\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nimport tensorflow as tf\nimport matplotlib.pyplot as plt\nimport xml.etree.ElementTree as ET\n\nfrom PIL import Image\nfrom scipy import linalg\nfrom sklearn.metrics import *\nfrom collections import Counter\nfrom scipy.stats import truncnorm\nfrom tqdm import tqdm_notebook as tqdm\nfrom multiprocessing import Pool, cpu_count\nfrom sklearn.preprocessing import LabelEncoder\n\nimport torch\nimport torch as th\nimport torch.nn as nn\nimport torch.utils.data\nimport torchvision as tv\nimport torch.nn.functional as F\nimport torchvision.models as models\n\nfrom torch.optim import Adam\nfrom torch.nn import Parameter\n\nfrom torchvision.datasets import *\nfrom torch.autograd import Variable\nfrom torch.optim.lr_scheduler import *\nfrom torchvision.utils import save_image\nfrom torchvision.datasets.folder import *\nfrom torch.nn.functional import interpolate\nfrom torchvision.transforms import transforms\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.nn import Conv2d, BCEWithLogitsLoss, DataParallel, AvgPool2d, ModuleList, LeakyReLU, ConvTranspose2d, Embedding\n\nbegin = time.time()\nsns.set_style('white')\nwarnings.simplefilter(action='ignore', category=FutureWarning) \nwarnings.simplefilter(action='ignore', category=UserWarning)\nwarnings.simplefilter(action='ignore', category=RuntimeWarning)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Initialization"},{"metadata":{},"cell_type":"markdown","source":"## Load Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"KERNEL_START_TIME = time.time()","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"DATA_PATH = '../input/all-dogs/all-dogs/'\nIMAGE_FOLDER = '../input/all-dogs'\nIMAGES = [DATA_PATH + p for p in os.listdir(DATA_PATH)]\nANNOTATION_PATH = '../input/annotation/Annotation'\n\nprint('Number of doge images :', len(IMAGES))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"img = Image.open(IMAGES[0])\nplt.imshow(img)\nplt.title('Woof', size=15)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Seeding\n\nAlthough I used the usual seeding rooting, that is supposed to ensure reproductible results, there is still some randomness in this kernel. Not sure why."},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def seed_everything(seed):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n#     torch.backends.cudnn.benchmark = False","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"seed = 2019\nseed_everything(seed)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Datasets"},{"metadata":{},"cell_type":"markdown","source":"## Transforms\n- RandomHorizontalFlip, p=0.5\n- CenterCrop or RandomCrop\n- A bit of ColorJitter\n- Rescale between [-1, 1]\n\nRotating is not used, and I also tried other normalization methods."},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class Rescale():\n    def __init__(self):\n        self.a = 2\n        self.b = -1\n        \n    def __call__(self, tensor):\n        return tensor.mul(self.a).add(self.b)\n\n    def __repr__(self):\n        return self.__class__.__name__ + '(x{}, +{})'.format(self.a, self.b)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_transforms(size):\n    base_transforms = transforms.Compose([transforms.Resize(size)])\n    additional_transforms = transforms.Compose([\n        transforms.RandomHorizontalFlip(p=0.5),\n#         transforms.RandomApply([transforms.RandomRotation(degrees=5)], p=0.5),\n        transforms.RandomChoice([transforms.CenterCrop(size), transforms.RandomCrop(size)]),\n        transforms.RandomApply([transforms.ColorJitter(brightness=0.2, contrast=(0.9, 1.2), saturation=0.3, hue=0.01)], p=0.5),\n        transforms.ToTensor(),\n#         transforms.Normalize(means, stds),\n        Rescale(),\n    ])\n    return base_transforms, additional_transforms","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Annotations Dataset\n> Thanks to Guillaume Desforges\n\n- I used dog boxes, but enlarge them by 10px on each side\n- Dog races are also used"},{"metadata":{"trusted":true},"cell_type":"code","source":"class DogeDataset(Dataset):\n    def __init__(self, folder, base_transforms, additional_transforms):\n        self.folder = folder\n        self.classes = [dirname[10:] for dirname in os.listdir(ANNOTATION_PATH)]\n        \n        self.base_transforms = base_transforms\n        self.additional_transforms = additional_transforms\n        self.imgs, self.labels = self.load_subfolders_images(folder)\n        \n        le = LabelEncoder().fit(self.classes)\n        self.y = torch.from_numpy(le.transform(self.labels)).long()\n        self.classes = le.inverse_transform(range(len(self.classes)))\n        \n    def __getitem__(self, index):\n        return self.additional_transforms(self.imgs[index]), self.y[index]\n\n    def __len__(self):\n        return len(self.imgs)\n    \n    @staticmethod\n    def is_valid_file(x):\n        img_extensions = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')\n        return has_file_allowed_extension(x, img_extensions)\n    \n    @staticmethod\n    def get_bbox(o):\n        bndbox = o.find('bndbox')\n        xmin = int(bndbox.find('xmin').text)\n        ymin = int(bndbox.find('ymin').text)\n        xmax = int(bndbox.find('xmax').text)\n        ymax = int(bndbox.find('ymax').text)\n        return xmin, ymin, xmax, ymax\n    \n    @staticmethod\n    def larger_bbox(bbox, ximg, yimg, a=10):\n        xmin, ymin, xmax, ymax = bbox\n        xmin = max(xmin - a, 0)\n        ymin = max(ymin - a, 0)\n        xmax = min(xmax + a, ximg)\n        ymax = min(ymax + a, yimg)\n        return (xmin, ymin, xmax, ymax)\n    \n    def load_subfolders_images(self, root):\n        imgs = []\n        paths = []\n        labels = []\n        \n        for root, _, fnames in sorted(os.walk(root)):\n            for fname in sorted(fnames):\n                path = os.path.join(root, fname)\n                if self.is_valid_file(path):\n                    paths.append(path)\n\n        for path in paths:\n            img = default_loader(path)\n\n            annotation_basename = os.path.splitext(os.path.basename(path))[0]\n            annotation_dirname = next(dirname for dirname in os.listdir(ANNOTATION_PATH) if dirname.startswith(annotation_basename.split('_')[0]))\n            annotation_filename = os.path.join(ANNOTATION_PATH, annotation_dirname, annotation_basename)\n            label = annotation_dirname[10:]\n            tree = ET.parse(annotation_filename)\n            root = tree.getroot()\n            objects = root.findall('object')\n            for o in objects:\n                bbox = self.get_bbox(o)\n                bbox = self.larger_bbox(bbox, img.size[0], img.size[1])\n                object_img = self.base_transforms(img.crop(bbox))\n                imgs.append(object_img)\n                labels.append(label)\n        return imgs, labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nbase_transforms, additional_transforms = get_transforms(64)\ndataset = DogeDataset(DATA_PATH, base_transforms, additional_transforms)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"nb_classes = len(dataset.classes)\nprint(f'Number of classes : {nb_classes}')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"nb_dogs = len(dataset)\nprint(f'Number of dogs : {nb_dogs}')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Metric\nNot used for submission"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"WEIGHTS_PATH = '../input/dog-face-generation-competition-kid-metric-input/classify_image_graph_def.pb'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"model_params = {\n    'Inception': {\n        'name': 'Inception', \n        'imsize': 64,\n        'output_layer': 'Pretrained_Net/pool_3:0', \n        'input_layer': 'Pretrained_Net/ExpandDims:0',\n        'output_shape': 2048,\n        'cosine_distance_eps': 0.1\n        }\n}","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class KernelEvalException(Exception):\n    pass","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"def create_model_graph(pth):\n    with tf.gfile.FastGFile( pth, 'rb') as f:\n        graph_def = tf.GraphDef()\n        graph_def.ParseFromString( f.read())\n        _ = tf.import_graph_def( graph_def, name='Pretrained_Net')\n\n        \ndef _get_model_layer(sess, model_name):\n    layername = model_params[model_name]['output_layer']\n    layer = sess.graph.get_tensor_by_name(layername)\n    ops = layer.graph.get_operations()\n    for op_idx, op in enumerate(ops):\n        for o in op.outputs:\n            shape = o.get_shape()\n            if shape._dims != []:\n                shape = [s.value for s in shape]\n                new_shape = []\n                for j, s in enumerate(shape):\n                    if s == 1 and j == 0:\n                        new_shape.append(None)\n                    else:\n                        new_shape.append(s)\n                o.__dict__['_shape_val'] = tf.TensorShape(new_shape)\n    return layer\n\n\ndef get_activations(images, sess, model_name, batch_size=50, verbose=False):\n    inception_layer = _get_model_layer(sess, model_name)\n    n_images = images.shape[0]\n    if batch_size > n_images:\n        print(\"warning: batch size is bigger than the data size. setting batch size to data size\")\n        batch_size = n_images\n    n_batches = n_images//batch_size + 1\n    pred_arr = np.empty((n_images,model_params[model_name]['output_shape']))\n    for i in range(n_batches):\n        if verbose:\n            print(\"\\rPropagating batch %d/%d\" % (i+1, n_batches), end=\"\", flush=True)\n        start = i*batch_size\n        if start+batch_size < n_images:\n            end = start+batch_size\n        else:\n            end = n_images\n                    \n        batch = images[start:end]\n        pred = sess.run(inception_layer, {model_params[model_name]['input_layer']: batch})\n        pred_arr[start:end] = pred.reshape(-1,model_params[model_name]['output_shape'])\n    if verbose:\n        print(\" done\")\n    return pred_arr","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def normalize_rows(x: np.ndarray):\n    return np.nan_to_num(x/np.linalg.norm(x, ord=2, axis=1, keepdims=True))\n\n\ndef cosine_distance(features1, features2):\n    features1_nozero = features1[np.sum(features1, axis=1) != 0]\n    features2_nozero = features2[np.sum(features2, axis=1) != 0]\n    norm_f1 = normalize_rows(features1_nozero)\n    norm_f2 = normalize_rows(features2_nozero)\n\n    d = 1.0 - np.abs(np.matmul(norm_f1, norm_f2.T))\n    mean_min_d = np.mean(np.min(d, axis=1))\n    return mean_min_d\n\n\ndef distance_thresholding(d, eps):\n    if d < eps:\n        return d\n    else:\n        return 1\n\n    \ndef calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):\n    mu1 = np.atleast_1d(mu1)\n    mu2 = np.atleast_1d(mu2)\n\n    sigma1 = np.atleast_2d(sigma1)\n    sigma2 = np.atleast_2d(sigma2)\n\n    assert mu1.shape == mu2.shape, \"Training and test mean vectors have different lengths\"\n    assert sigma1.shape == sigma2.shape, \"Training and test covariances have different dimensions\"\n\n    diff = mu1 - mu2\n\n    # product might be almost singular\n    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)\n    if not np.isfinite(covmean).all():\n        msg = \"fid calculation produces singular product; adding %s to diagonal of cov estimates\" % eps\n        warnings.warn(msg)\n        offset = np.eye(sigma1.shape[0]) * eps\n        # covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))\n        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))\n    \n    # numerical error might give slight imaginary component\n    if np.iscomplexobj(covmean):\n        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):\n            m = np.max(np.abs(covmean.imag))\n            raise ValueError(\"Imaginary component {}\".format(m))\n        covmean = covmean.real\n\n    tr_covmean = np.trace(covmean)\n    return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean\n\n\ndef calculate_activation_statistics(images, sess, model_name, batch_size=50, verbose=False):\n    act = get_activations(images, sess, model_name, batch_size, verbose)\n    mu = np.mean(act, axis=0)\n    sigma = np.cov(act, rowvar=False)\n    return mu, sigma, act\n\n\ndef _handle_path_memorization(path, sess, model_name, is_checksize, is_check_png):\n    path = pathlib.Path(path)\n    files = list(path.glob('*.jpg')) + list(path.glob('*.png'))\n    imsize = model_params[model_name]['imsize']\n\n    # In production we don't resize input images. This is just for demo purpose. \n    x = np.array([np.array(img_read_checks(fn, imsize, is_checksize, imsize, is_check_png)) for fn in files])\n    m, s, features = calculate_activation_statistics(x, sess, model_name)\n    del x #clean up memory\n    return m, s, features\n\n\ndef img_read_checks(filename, resize_to, is_checksize=False, check_imsize = 64, is_check_png = False):\n    im = Image.open(str(filename))\n    if is_checksize and im.size != (check_imsize,check_imsize):\n        raise KernelEvalException('The images are not of size '+str(check_imsize))\n    \n    if is_check_png and im.format != 'PNG':\n        raise KernelEvalException('Only PNG images should be submitted.')\n\n    if resize_to is None:\n        return im\n    else:\n        return im.resize((resize_to,resize_to),Image.ANTIALIAS)\n\n    \ndef calculate_kid_given_paths(paths, model_name, model_path, feature_path=None):\n    ''' Calculates the KID of two paths. '''\n    tf.reset_default_graph()\n    create_model_graph(str(model_path))\n    with tf.Session() as sess:\n        sess.run(tf.global_variables_initializer())\n        m1, s1, features1 = _handle_path_memorization(paths[0], sess, model_name, is_checksize = True, is_check_png = True)\n        if feature_path is None:\n            m2, s2, features2 = _handle_path_memorization(paths[1], sess, model_name, is_checksize = False, is_check_png = False)\n        else:\n            with np.load(feature_path) as f:\n                m2, s2, features2 = f['m'], f['s'], f['features']\n\n#         print('m1,m2 shape=',(m1.shape,m2.shape),'s1,s2=',(s1.shape,s2.shape))\n#         print('starting calculating FID')\n        fid_value = calculate_frechet_distance(m1, s1, m2, s2)\n#         print('done with FID, starting distance calculation')\n        distance = cosine_distance(features1, features2)        \n        return fid_value, distance","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def compute_mifid(generated_path, real_path, graph_path, model_params, eps=10e-15):\n    fid_value, distance = calculate_kid_given_paths([generated_path, real_path], \n                                                                  'Inception', graph_path)\n    distance = distance_thresholding(distance, model_params['Inception']['cosine_distance_eps'])\n    return fid_value / (distance + eps)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Models"},{"metadata":{},"cell_type":"markdown","source":"## Tools"},{"metadata":{},"cell_type":"markdown","source":"### Spectral Normalization\n> Taken from https://github.com/isr-wang/SNGAN\n\nI played a bit with it, without success."},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def l2normalize(v, eps=1e-12):\n    return v / (v.norm() + eps)\n\n\nclass SpectralNorm(nn.Module):\n    def __init__(self, module, name='weight', power_iterations=1):\n        super(SpectralNorm, self).__init__()\n        self.module = module\n        self.name = name\n        self.power_iterations = power_iterations\n        if not self._made_params():\n            self._make_params()\n\n    def _update_u_v(self):\n        u = getattr(self.module, self.name + \"_u\")\n        v = getattr(self.module, self.name + \"_v\")\n        w = getattr(self.module, self.name + \"_bar\")\n\n        height = w.data.shape[0]\n        for _ in range(self.power_iterations):\n            v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data))\n            u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data))\n\n        # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data))\n        sigma = u.dot(w.view(height, -1).mv(v))\n        setattr(self.module, self.name, w / sigma.expand_as(w))\n\n    def _made_params(self):\n        try:\n            u = getattr(self.module, self.name + \"_u\")\n            v = getattr(self.module, self.name + \"_v\")\n            w = getattr(self.module, self.name + \"_bar\")\n            return True\n        except AttributeError:\n            return False\n\n\n    def _make_params(self):\n        w = getattr(self.module, self.name)\n\n        height = w.data.shape[0]\n        width = w.view(height, -1).data.shape[1]\n\n        u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False)\n        v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False)\n        u.data = l2normalize(u.data)\n        v.data = l2normalize(v.data)\n        w_bar = Parameter(w.data)\n\n        del self.module._parameters[self.name]\n\n        self.module.register_parameter(self.name + \"_u\", u)\n        self.module.register_parameter(self.name + \"_v\", v)\n        self.module.register_parameter(self.name + \"_bar\", w_bar)\n\n\n    def forward(self, *args):\n        self._update_u_v()\n        return self.module.forward(*args)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Equalized Layers"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"# extending Conv2D and Deconv2D layers for equalized learning rate logic\nclass _equalized_conv2d(th.nn.Module):\n    \"\"\" conv2d with the concept of equalized learning rate\n        Args:\n            :param c_in: input channels\n            :param c_out:  output channels\n            :param k_size: kernel size (h, w) should be a tuple or a single integer\n            :param stride: stride for conv\n            :param pad: padding\n            :param bias: whether to use bias or not\n    \"\"\"\n\n    def __init__(self, c_in, c_out, k_size, stride=1, pad=0, bias=True):\n        \"\"\" constructor for the class \"\"\"\n        from torch.nn.modules.utils import _pair\n        from numpy import sqrt, prod\n\n        super(_equalized_conv2d, self).__init__()\n\n        # define the weight and bias if to be used\n        self.weight = th.nn.Parameter(th.nn.init.normal_(\n            th.empty(c_out, c_in, *_pair(k_size))\n        ))\n\n        self.use_bias = bias\n        self.stride = stride\n        self.pad = pad\n\n        if self.use_bias:\n            self.bias = th.nn.Parameter(th.FloatTensor(c_out).fill_(0))\n\n        fan_in = prod(_pair(k_size)) * c_in  # value of fan_in\n        self.scale = sqrt(2) / sqrt(fan_in)\n\n    def forward(self, x):\n        \"\"\"\n        forward pass of the network\n        :param x: input\n        :return: y => output\n        \"\"\"\n        from torch.nn.functional import conv2d\n\n        return conv2d(input=x,\n                      weight=self.weight * self.scale,  # scale the weight on runtime\n                      bias=self.bias if self.use_bias else None,\n                      stride=self.stride,\n                      padding=self.pad)\n\n    def extra_repr(self):\n        return \", \".join(map(str, self.weight.shape))\n\n\nclass _equalized_deconv2d(th.nn.Module):\n    \"\"\" Transpose convolution using the equalized learning rate\n        Args:\n            :param c_in: input channels\n            :param c_out: output channels\n            :param k_size: kernel size\n            :param stride: stride for convolution transpose\n            :param pad: padding\n            :param bias: whether to use bias or not\n    \"\"\"\n\n    def __init__(self, c_in, c_out, k_size, stride=1, pad=0, bias=True):\n        \"\"\" constructor for the class \"\"\"\n        from torch.nn.modules.utils import _pair\n        from numpy import sqrt\n\n        super(_equalized_deconv2d, self).__init__()\n\n        # define the weight and bias if to be used\n        self.weight = th.nn.Parameter(th.nn.init.normal_(\n            th.empty(c_in, c_out, *_pair(k_size))\n        ))\n\n        self.use_bias = bias\n        self.stride = stride\n        self.pad = pad\n\n        if self.use_bias:\n            self.bias = th.nn.Parameter(th.FloatTensor(c_out).fill_(0))\n\n        fan_in = c_in  # value of fan_in for deconv\n        self.scale = sqrt(2) / sqrt(fan_in)\n\n    def forward(self, x):\n        \"\"\"\n        forward pass of the layer\n        :param x: input\n        :return: y => output\n        \"\"\"\n        from torch.nn.functional import conv_transpose2d\n\n        return conv_transpose2d(input=x,\n                                weight=self.weight * self.scale,  # scale the weight on runtime\n                                bias=self.bias if self.use_bias else None,\n                                stride=self.stride,\n                                padding=self.pad)\n\n    def extra_repr(self):\n        return \", \".join(map(str, self.weight.shape))\n\n\nclass _equalized_linear(th.nn.Module):\n    \"\"\" Linear layer using equalized learning rate\n        Args:\n            :param c_in: number of input channels\n            :param c_out: number of output channels\n            :param bias: whether to use bias with the linear layer\n    \"\"\"\n\n    def __init__(self, c_in, c_out, bias=True):\n        \"\"\"\n        Linear layer modified for equalized learning rate\n        \"\"\"\n        from numpy import sqrt\n\n        super(_equalized_linear, self).__init__()\n\n        self.weight = th.nn.Parameter(th.nn.init.normal_(\n            th.empty(c_out, c_in)\n        ))\n\n        self.use_bias = bias\n\n        if self.use_bias:\n            self.bias = th.nn.Parameter(th.FloatTensor(c_out).fill_(0))\n\n        fan_in = c_in\n        self.scale = sqrt(2) / sqrt(fan_in)\n\n    def forward(self, x):\n        \"\"\"\n        forward pass of the layer\n        :param x: input\n        :return: y => output\n        \"\"\"\n        from torch.nn.functional import linear\n        return linear(x, self.weight * self.scale,\n                      self.bias if self.use_bias else None)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Pixelwise Normalization"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class PixelwiseNorm(th.nn.Module):\n    def __init__(self):\n        super(PixelwiseNorm, self).__init__()\n\n    def forward(self, x, alpha=1e-8):\n        \"\"\"\n        forward pass of the module\n        :param x: input activations volume\n        :param alpha: small number for numerical stability\n        :return: y => pixel normalized activations\n        \"\"\"\n        y = x.pow(2.).mean(dim=1, keepdim=True).add(alpha).sqrt()  # [N1HW]\n        y = x / y  # normalize the input x volume\n        return y","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Minibatch Standard Deviation"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class MinibatchStdDev(th.nn.Module):\n    \"\"\"\n    Minibatch standard deviation layer for the discriminator\n    \"\"\"\n    def __init__(self):\n        \"\"\"\n        derived class constructor\n        \"\"\"\n        super(MinibatchStdDev, self).__init__()\n\n    def forward(self, x, alpha=1e-8):\n        \"\"\"\n        forward pass of the layer\n        :param x: input activation volume\n        :param alpha: small number for numerical stability\n        :return: y => x appended with standard deviation constant map\n        \"\"\"\n        batch_size, _, height, width = x.shape\n\n        # [B x C x H x W] Subtract mean over batch.\n        y = x - x.mean(dim=0, keepdim=True)\n\n        # [1 x C x H x W]  Calc standard deviation over batch\n        y = th.sqrt(y.pow(2.).mean(dim=0, keepdim=False) + alpha)\n\n        # [1]  Take average over feature_maps and pixels.\n        y = y.mean().view(1, 1, 1, 1)\n\n        # [B x 1 x H x W]  Replicate over group and pixels.\n        y = y.repeat(batch_size, 1, height, width)\n\n        # [B x C x H x W]  Append as new feature_map.\n        y = th.cat([x, y], 1)\n\n        # return the computed values:\n        return y","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Generator Blocks"},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class GenInitialBlock(th.nn.Module):\n    def __init__(self, in_channels, use_eql, use_spec_norm=False):\n        super(GenInitialBlock, self).__init__()\n\n        if use_eql:\n            self.conv_1 = _equalized_deconv2d(in_channels, in_channels, (4, 4), bias=True)\n            self.conv_2 = _equalized_conv2d(in_channels, in_channels, (3, 3), pad=1, bias=True)\n\n        else:\n            self.conv_1 = ConvTranspose2d(in_channels, in_channels, (4, 4), bias=True)\n            self.conv_2 = Conv2d(in_channels, in_channels, (3, 3), padding=1, bias=True)\n\n        if use_spec_norm:\n            self.conv_1 = SpectralNorm(self.conv_1)\n            self.conv_2 = SpectralNorm(self.conv_2)\n            \n        self.pixNorm = PixelwiseNorm()\n        self.lrelu = LeakyReLU(0.2)\n        \n    def forward(self, x):\n        # convert the tensor shape:\n        y = th.unsqueeze(th.unsqueeze(x, -1), -1)\n\n        # perform the forward computations:\n        y = self.lrelu(self.conv_1(y))\n        y = self.lrelu(self.conv_2(y))\n        y = self.pixNorm(y)\n        \n        return y","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class GenGeneralConvBlock(th.nn.Module):\n    def __init__(self, in_channels, out_channels, use_eql, use_spec_norm=False):\n        super(GenGeneralConvBlock, self).__init__()\n\n        self.upsample = lambda x: interpolate(x, scale_factor=2)\n\n        if use_eql:\n            self.conv_1 = _equalized_conv2d(in_channels, out_channels, (3, 3),\n                                            pad=1, bias=True)\n            self.conv_2 = _equalized_conv2d(out_channels, out_channels, (3, 3),\n                                            pad=1, bias=True)\n        else:\n            self.conv_1 = Conv2d(in_channels, out_channels, (3, 3),\n                                 padding=1, bias=True)\n            self.conv_2 = Conv2d(out_channels, out_channels, (3, 3),\n                                 padding=1, bias=True)\n        \n        if use_spec_norm:\n            self.conv_1 = SpectralNorm(self.conv_1)\n            self.conv_2 = SpectralNorm(self.conv_2)\n            \n        self.pixNorm = PixelwiseNorm()\n        self.lrelu = LeakyReLU(0.2)\n\n    def forward(self, x):\n        \"\"\"\n        forward pass of the block\n        :param x: input\n        :return: y => output\n        \"\"\"\n        y = self.upsample(x)\n        y = self.pixNorm(self.lrelu(self.conv_1(y)))\n        y = self.pixNorm(self.lrelu(self.conv_2(y)))\n        \n        return y","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Discriminator Blocks"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class DisGeneralConvBlock(th.nn.Module):\n    \"\"\" General block in the discriminator  \"\"\"\n\n    def __init__(self, in_channels, out_channels, use_eql, use_spec_norm=False):\n        \"\"\"\n        constructor of the class\n        :param in_channels: number of input channels\n        :param out_channels: number of output channels\n        :param use_eql: whether to use equalized learning rate\n        \"\"\"\n        super(DisGeneralConvBlock, self).__init__()\n        \n        if use_eql:\n            self.conv_1 = _equalized_conv2d(in_channels, in_channels, (3, 3), pad=1, bias=True)\n            self.conv_2 = _equalized_conv2d(in_channels, out_channels, (3, 3), pad=1, bias=True)\n        else:\n            self.conv_1 = Conv2d(in_channels, in_channels, (3, 3), padding=1, bias=True)\n            self.conv_2 = Conv2d(in_channels, out_channels, (3, 3), padding=1, bias=True)\n            \n        if use_spec_norm:\n            self.conv_1 = SpectralNorm(self.conv_1)\n            self.conv_2 = SpectralNorm(self.conv_2)\n\n        self.downSampler = AvgPool2d(2)\n        self.lrelu = LeakyReLU(0.2)\n\n    def forward(self, x):\n        y = self.lrelu(self.conv_1(x))\n        y = self.lrelu(self.conv_2(y))\n        y = self.downSampler(y)\n        return y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class ConDisFinalBlock(th.nn.Module):\n    def __init__(self, in_channels, num_classes, use_eql):\n        super(ConDisFinalBlock, self).__init__()\n\n        # declare the required modules for forward pass\n        self.batch_discriminator = MinibatchStdDev()\n        if use_eql:\n            self.conv_1 = _equalized_conv2d(in_channels + 1, in_channels, (3, 3), pad=1, bias=True)\n            self.conv_2 = _equalized_conv2d(in_channels, in_channels, (4, 4), bias=True)\n\n            # final conv layer emulates a fully connected layer\n            self.conv_3 = _equalized_conv2d(in_channels, 1, (1, 1), bias=True)\n        else:\n            self.conv_1 = Conv2d(in_channels + 1, in_channels, (3, 3), padding=1, bias=True)\n            self.conv_2 = Conv2d(in_channels, in_channels, (4, 4), bias=True)\n\n            # final conv layer emulates a fully connected layer\n            self.conv_3 = Conv2d(in_channels, 1, (1, 1), bias=True)\n\n        # we also need an embedding matrix for the label vectors\n        self.label_embedder = Embedding(num_classes, in_channels, max_norm=1)\n\n        # leaky_relu:\n        self.lrelu = LeakyReLU(0.2)\n        \n        # for ft matching\n        nb_ft = 128\n        self.ft_matching_dense = nn.Linear(2 * in_channels, nb_ft)\n\n    def forward(self, x, labels, return_ft=False):\n        \"\"\"\n        forward pass of the FinalBlock\n        :param x: input\n        :param labels: samples' labels for conditional discrimination\n                       Note that these are pure integer labels [Batch_size x 1]\n        :return: y => output\n        \"\"\"\n        batch_size = x.size()[0]\n        # minibatch_std_dev layer\n        y = self.batch_discriminator(x)  # [B x C x 4 x 4]\n\n        # perform the forward pass\n        y = self.lrelu(self.conv_1(y))  # [B x C x 4 x 4]\n\n        # obtain the computed features\n        y = self.lrelu(self.conv_2(y))  # [B x C x 1 x 1]\n        y_ = y.view((batch_size, -1))  # [B x C]\n        \n        # embed the labels\n        labels = self.label_embedder(labels.cuda()).view((batch_size, -1))  # [B x C]\n\n        # compute the inner product with the label embeddings\n        \n        if return_ft:\n            self.ft_matching_dense(torch.cat((y_, labels), 1))\n        \n        projection_scores = (y_ * labels).sum(dim=-1)  # [B]\n        \n        # normal discrimination score\n        y = self.lrelu(self.conv_3(y))  # This layer has linear activation\n\n        # calculate the total score\n        final_score = y.view(-1) + projection_scores\n\n        # return the output raw discriminator scores\n        return final_score","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Gaussian Noise\nTried it as well, did not improve the results."},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class GaussianNoise(nn.Module):\n    def __init__(self, sigma=0.1):\n        super().__init__()\n        self.sigma = sigma\n        self.noise = torch.tensor(0).cuda()\n\n    def forward(self, x):\n        if self.training:\n            noise = self.noise.repeat(*x.size()).float().normal_() * self.sigma #self.sigma * x?\n            return x + noise\n        return x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Generator"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, depth=5, latent_size=128, use_eql=True, use_spec_norm=False):\n        super(Generator, self).__init__()\n\n        assert latent_size != 0 and ((latent_size & (latent_size - 1)) == 0), \"latent size not a power of 2\"\n        if depth >= 4: assert latent_size >= np.power(2, depth - 4), \"latent size will diminish to zero\"\n\n        self.use_eql = use_eql\n        self.use_spec_norm = use_spec_norm\n        self.depth = depth\n        self.latent_size = latent_size\n\n        # register the modules required for the GAN\n        self.initial_block = GenInitialBlock(self.latent_size, use_eql=self.use_eql, use_spec_norm=False)\n\n        # create a module list of the other required general convolution blocks\n        self.layers = ModuleList([])  # initialize to empty list\n\n        # create the ToRGB layers for various outputs:\n        if self.use_eql:\n            self.toRGB = lambda in_channels: _equalized_conv2d(in_channels, 3, (1, 1), bias=True)\n        else:\n            self.toRGB = lambda in_channels: Conv2d(in_channels, 3, (1, 1), bias=True)\n\n        self.rgb_converters = ModuleList([self.toRGB(self.latent_size)])\n\n        # create the remaining layers\n        for i in range(self.depth - 1):\n            if i <= 2:\n                layer = GenGeneralConvBlock(self.latent_size, self.latent_size, use_eql=self.use_eql, use_spec_norm=use_spec_norm)\n                rgb = self.toRGB(self.latent_size)\n            else:\n                in_size = int(self.latent_size // np.power(2, i - 3))\n                out_size = int(self.latent_size // np.power(2, i - 2))\n                \n                layer = nn.Sequential(\n                    GenGeneralConvBlock(in_size, out_size, use_eql=self.use_eql, use_spec_norm=use_spec_norm),\n#                     Self_Attn(out_size)\n                )\n                rgb = self.toRGB(out_size)\n                \n            self.layers.append(layer)\n            self.rgb_converters.append(rgb)\n\n        # register the temporary upsampler\n        self.temporaryUpsampler = lambda x: interpolate(x, scale_factor=2)\n        self.tanh = nn.Tanh()\n\n    def forward(self, x, depth, alpha):\n        assert depth < self.depth, \"Requested output depth cannot be produced\"\n\n        y = self.initial_block(x)\n\n        if depth > 0:\n            for block in self.layers[:depth - 1]:\n                y = block(y)\n\n            residual = self.rgb_converters[depth - 1](self.temporaryUpsampler(y))\n            straight = self.rgb_converters[depth](self.layers[depth - 1](y))\n\n            out = (alpha * straight) + ((1 - alpha) * residual)\n\n        else:\n            out = self.rgb_converters[0](y)\n\n        return self.tanh(out)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Discriminator"},{"metadata":{"trusted":true},"cell_type":"code","source":"class ConditionalDiscriminator(nn.Module):\n    def __init__(self, num_classes, height=7, feature_size=512, use_eql=True, use_spec_norm=False):\n        super(ConditionalDiscriminator, self).__init__()\n\n        assert feature_size != 0 and ((feature_size & (feature_size - 1)) == 0), \"latent size not a power of 2\"\n        if height >= 4:\n            assert feature_size >= np.power(2, height - 4), \"feature size cannot be produced\"\n\n        # create state of the object\n        self.use_eql = use_eql\n        self.use_spec_norm = use_spec_norm\n        self.height = height\n        self.feature_size = feature_size\n        self.num_classes = num_classes\n        \n        self.noise = GaussianNoise(sigma=0.2)\n        self.final_block = ConDisFinalBlock(self.feature_size, self.num_classes, use_eql=self.use_eql)\n\n        # create a module list of the other required general convolution blocks\n        self.layers = ModuleList([])  # initialize to empty list\n\n        # create the fromRGB layers for various inputs:\n        if self.use_eql:\n            self.fromRGB = lambda out_channels: _equalized_conv2d(3, out_channels, (1, 1), bias=True)\n        else:\n            self.fromRGB = lambda out_channels: Conv2d(3, out_channels, (1, 1), bias=True)\n\n        rgb = self.fromRGB(self.feature_size)\n        if use_spec_norm: \n            rgb = SpectralNorm(rgb)\n        self.rgb_to_features = ModuleList([rgb])\n\n        # create the remaining layers\n        for i in range(self.height - 1):\n            if i > 2:\n                in_size = int(self.feature_size // np.power(2, i - 2))\n                out_size = int(self.feature_size // np.power(2, i - 3))\n                layer = nn.Sequential(\n                    DisGeneralConvBlock(in_size, out_size, use_eql=self.use_eql, use_spec_norm=use_spec_norm),\n#                     self.noise,\n                )\n                \n                rgb = self.fromRGB(in_size)\n\n            else:\n                layer = nn.Sequential(\n                    DisGeneralConvBlock(self.feature_size, self.feature_size,  use_eql=self.use_eql, use_spec_norm=use_spec_norm),\n#                     self.noise,\n                )\n                rgb = self.fromRGB(self.feature_size)\n            \n            if use_spec_norm: \n                rgb = SpectralNorm(rgb)\n\n            self.layers.append(layer)\n            self.rgb_to_features.append(rgb)\n\n        # register the temporary downSampler\n        self.temporaryDownsampler = AvgPool2d(2)\n\n    def forward(self, x, labels, height, alpha, return_ft=False):\n        assert height < self.height, \"Requested output depth cannot be produced\"\n\n        if height > 0:\n            residual = self.rgb_to_features[height - 1](self.temporaryDownsampler(x))\n\n            straight = self.layers[height - 1](\n                self.rgb_to_features[height](x)\n            )\n\n            y = (alpha * straight) + ((1 - alpha) * residual)\n\n            for block in reversed(self.layers[:height - 1]):\n                y = block(y)\n        else:\n            y = self.rgb_to_features[0](x)\n\n        out = self.final_block(y, labels, return_ft=return_ft)\n        return out","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Tools"},{"metadata":{},"cell_type":"markdown","source":"## Losses\nI first went with the usual loss with smooth and noisy labels, but finally set up with the Hinge one."},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class ConditionalGANLoss:\n    \"\"\" Base class for all conditional losses \"\"\"\n    def __init__(self, dis):\n        self.dis = dis\n\n    def dis_loss(self, real_samps, fake_samps, labels, height, alpha):\n        raise NotImplementedError(\"dis_loss method has not been implemented\")\n\n    def gen_loss(self, real_samps, fake_samps, labels, height, alpha):\n        raise NotImplementedError(\"gen_loss method has not been implemented\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"class StandardLoss(ConditionalGANLoss):\n    def __init__(self, dis):\n        super().__init__(dis)\n        self.criterion = BCEWithLogitsLoss(reduction='sum')\n\n    def dis_loss(self, real_samps, fake_samps, labels, height, alpha):\n        assert real_samps.device == fake_samps.device, \"Different devices\"\n        \n        preds_real = self.dis(real_samps, labels, height, alpha)\n        preds_fake = self.dis(fake_samps, labels, height, alpha)\n        \n        labels_real = torch.from_numpy(np.random.uniform(0.5, 0.99, real_samps.size()[0])).float().cuda()\n        labels_fake = torch.from_numpy(np.random.uniform(0, 0.25, fake_samps.size()[0])).float().cuda()\n\n        real_loss = self.criterion(preds_real.view(-1), labels_real)\n        fake_loss = self.criterion(preds_fake.view(-1), labels_fake)\n\n        return real_loss + fake_loss\n\n    def gen_loss(self, _, fake_samps, labels, height, alpha):\n        preds_fake = self.dis(fake_samps, labels, height, alpha)\n        labels_real = torch.from_numpy(np.random.uniform(0.5, 0.99, fake_samps.size()[0])).float().cuda()\n        return self.criterion(preds_fake.view(-1), labels_real)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Hinge(ConditionalGANLoss):\n    def __init__(self, dis):\n        super().__init__(dis)\n\n    def dis_loss(self, real_samps, fake_samps, labels, height, alpha):\n        r_preds = self.dis(real_samps, labels, height, alpha)\n        f_preds = self.dis(fake_samps, labels, height, alpha)\n        loss = (torch.mean(th.nn.ReLU()(1 - r_preds)) + torch.mean(th.nn.ReLU()(1 + f_preds)))\n        return loss\n\n    def gen_loss(self, _, fake_samps, labels, height, alpha):\n        return - torch.mean(self.dis(fake_samps, labels, height, alpha))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## EMA Weights Decay\nTo stabilize the training of the generator"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"def update_average(model_old, model_new, beta):\n    def toggle_grad(model, requires_grad):\n        for p in model.parameters():\n            p.requires_grad_(requires_grad)\n\n    toggle_grad(model_old, False)\n    toggle_grad(model_new, False)\n\n    param_dict_new = dict(model_new.named_parameters())\n    for param_name, param_old in model_old.named_parameters():\n        param_new = param_dict_new[param_name]\n        assert (param_old is not param_new)\n        param_old.copy_(beta * param_old + (1. - beta) * param_new)\n\n    toggle_grad(model_old, True)\n    toggle_grad(model_new, True)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Save & Load Models"},{"metadata":{"trusted":true},"cell_type":"code","source":"def save_model_weights(model, filename, verbose=1):\n    if verbose:\n        print(f'-> Saving weights to {filename}')\n    torch.save(model.state_dict(), filename)\n    \ndef load_model_weights(model, filename, verbose=1):\n    if verbose:\n        print(f'-> Loading weights from {filename}')\n    model.load_state_dict(torch.load(filename))\n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Conditional ProGAN Wrapper\nA big piece of PyTorch code that deserves some cleaning."},{"metadata":{"trusted":true},"cell_type":"code","source":"class ConditionalProGAN:\n    def __init__(self, num_classes=120, depth=7, latent_size=128, embed_dim=64,\n                 lr_g=0.001, lr_d=0.001, n_critic=1, use_eql=True, use_spec_norm=False,\n                 loss=StandardLoss, use_ema=True, ema_decay=0.999):\n\n        self.gen = Generator(depth=depth, latent_size=latent_size, \n                             use_eql=use_eql, use_spec_norm=False).cuda()\n        self.dis = ConditionalDiscriminator(num_classes, height=depth, feature_size=latent_size, \n                                            use_eql=use_eql, use_spec_norm=use_spec_norm).cuda()\n        \n        self.gen = DataParallel(self.gen)\n        self.dis = DataParallel(self.dis)\n\n        self.latent_size = latent_size\n        self.num_classes = num_classes\n        self.depth = depth\n        \n        self.use_ema = use_ema\n        self.ema_decay = ema_decay\n        self.n_critic = n_critic\n        self.use_eql = use_eql\n        self.drift = 0.001\n        \n        self.lr_g = lr_g\n        self.lr_d = lr_d\n        \n        self.gen_optim = Adam(self.gen.parameters(), lr=self.lr_g, betas=(0.5, 0.99), eps=1e-8)\n        self.dis_optim = Adam(self.dis.parameters(), lr=self.lr_d, betas=(0.5, 0.99), eps=1e-8)\n\n        try:\n            self.loss = loss(self.dis)\n        except:\n            self.loss = loss(self.dis, drift=self.drift, use_gp=True)\n\n        # setup the ema for the generator\n        if self.use_ema:\n            self.gen_shadow = copy.deepcopy(self.gen)\n            self.ema_updater = update_average\n            self.ema_updater(self.gen_shadow, self.gen, beta=0)\n\n    \n    def __progressive_downsampling(self, real_batch, depth, alpha):\n        \"\"\"\n        private helper for downsampling the original images in order to facilitate the\n        progressive growing of the layers.\n        :param real_batch: batch of real samples\n        :param depth: depth at which training is going on\n        :param alpha: current value of the fader alpha\n        :return: real_samples => modified real batch of samples\n        \"\"\"\n        \n        # downsample the real_batch for the given depth\n        down_sample_factor = int(np.power(2, self.depth - depth - 1))\n        prior_downsample_factor = max(int(np.power(2, self.depth - depth)), 0)\n\n        ds_real_samples = AvgPool2d(down_sample_factor)(real_batch)\n\n        if depth > 0:\n            prior_ds_real_samples = interpolate(AvgPool2d(prior_downsample_factor)(real_batch),\n                                                scale_factor=2)\n        else:\n            prior_ds_real_samples = ds_real_samples\n\n        # real samples are a combination of ds_real_samples and prior_ds_real_samples\n        real_samples = (alpha * ds_real_samples) + ((1 - alpha) * prior_ds_real_samples)\n\n        # return the so computed real_samples\n        return real_samples\n\n    \n    def optimize_discriminator(self, noise, real_batch, labels, depth, alpha):\n        real_samples = self.__progressive_downsampling(real_batch, depth, alpha)\n        loss_val = 0\n        \n        for _ in range(self.n_critic):\n            fake_samples = self.gen(noise, depth, alpha).detach()\n            loss = self.loss.dis_loss(real_samples, fake_samples, labels, depth, alpha)\n\n            self.dis_optim.zero_grad()\n            loss.backward()\n            self.dis_optim.step()\n\n            loss_val += loss.item()\n\n        return loss_val / self.n_critic\n\n    \n    def optimize_generator(self, noise, real_batch, labels, depth, alpha):\n        real_samples = self.__progressive_downsampling(real_batch, depth, alpha)\n        fake_samples = self.gen(noise, depth, alpha)\n\n        loss = self.loss.gen_loss(real_samples, fake_samples, labels, depth, alpha)\n\n        self.gen_optim.zero_grad()\n        loss.backward()\n        self.gen_optim.step()\n\n        if self.use_ema:\n            self.ema_updater(self.gen_shadow, self.gen, self.ema_decay)\n\n        return loss.item()\n        \n        \n    def one_hot_encode(self, labels):\n        if not hasattr(self, \"label_oh_encoder\"):\n            self.label_oh_encoder = th.nn.Embedding(self.num_classes, self.num_classes)\n            self.label_oh_encoder.weight.data = th.eye(self.num_classes)\n        return self.label_oh_encoder(labels.view(-1))\n    \n    \n    @staticmethod\n    def scale(imgs):\n        def norm(img, inf, sup):\n            img.clamp_(min=inf, max=sup)\n            img.add_(-inf).div_(sup - inf + 1e-5)   \n            \n        for img in imgs:\n            norm(img, float(img.min()), float(img.max()))\n        # ndarr = grid.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()\n    \n    \n    @staticmethod\n    def truncated_normal(size, threshold=1):\n        values = truncnorm.rvs(-threshold, threshold, size=size)\n        return values\n    \n    \n    def generate(self, depth=None, alpha=1, noise=None, races=None, n=64, n_plot=0):\n        if depth is None:\n            depth = self.depth - 1\n        if noise is None:\n            noise = th.randn(n, self.latent_size - self.num_classes).cuda()\n#             z = self.truncated_normal(size=(n, self.latent_size - self.num_classes))\n#             noise = torch.from_numpy(z).float().cuda()\n        if races is None:\n            races = torch.from_numpy(np.random.choice(range(self.num_classes), size=n)).long()\n        \n        label_information = self.one_hot_encode(races).cuda()\n        gan_input = th.cat((label_information, noise), dim=-1)\n        \n        if self.use_ema:\n            generated_images = self.gen_shadow(gan_input, depth, alpha).detach().cpu()\n        else:\n            generated_images = self.gen(gan_input, depth, alpha).detach().cpu()\n\n#         self.scale(generated_images)\n        generated_images.add_(1).div_(2)\n        images = generated_images.clone().numpy().transpose(0, 2, 3, 1)                 \n        \n        if n_plot >= 5:\n            plt.figure(figsize=(15, 3 * n_plot//5))\n            for i in range(n_plot):\n                plt.subplot(n_plot//5, 5, i+1)\n                plt.imshow(images[i])\n                plt.axis('off')\n                plt.title(dataset.classes[races.cpu().numpy()[i]])\n            plt.show()\n        return generated_images\n    \n    \n    def generate_score(self, depth=None, alpha=1, noise=None, races=None, n=64, n_plot=0):\n        if depth is None:\n            depth = self.depth - 1\n        if noise is None:\n            noise = th.randn(n, self.latent_size - self.num_classes).cuda()\n        if races is None:\n            races = torch.from_numpy(np.random.choice(range(self.num_classes), size=n)).long()\n        \n        label_information = self.one_hot_encode(races).cuda()\n        gan_input = th.cat((label_information, noise), dim=-1)\n        \n        if self.use_ema:\n            generated_images = self.gen_shadow(gan_input, depth, alpha).detach().cpu()\n        else:\n            generated_images = self.gen(gan_input, depth, alpha).detach().cpu()\n\n        generated_images.add_(1).div_(2)\n        images = generated_images.clone().numpy().transpose(0, 2, 3, 1)\n        scores = nn.Sigmoid()(self.dis(generated_images, races, depth, alpha)).cpu().detach().numpy()\n        \n        if n_plot >= 5:\n            plt.figure(figsize=(15, 3 * n_plot//5))\n            for i in range(n_plot):\n                plt.subplot(n_plot//5, 5, i+1)\n                plt.imshow(images[i])\n                plt.axis('off')\n                plt.title(dataset.classes[races.cpu().numpy()[i]] + f' - {scores[i]:.3f}' )\n            plt.show()\n            \n        return images, generated_images, scores, races.cpu().numpy()\n    \n    \n    def plot_race(self, race_idx, depth=4, alpha=1, n_plot=5, n=128):\n        races = np.concatenate((np.array([race_idx]* n_plot), \n                                np.random.choice(range(self.num_classes), size=n - n_plot)))\n  \n        races = torch.from_numpy(races).long()\n        self.generate(depth, alpha=alpha, races=races, n=n, n_plot=n_plot)\n    \n    \n    def compute_mifid(self, alpha=1, folder='../tmp_images', n_images=10000, im_batch_size=100):\n        if os.path.exists(folder):\n            shutil.rmtree(folder, ignore_errors=True)\n        os.mkdir(folder)\n\n        for i_b in range(0, n_images, im_batch_size):\n            gen_images = self.generate(n=im_batch_size)\n            for i_img in range(gen_images.size(0)):\n                save_image(gen_images[i_img, :, :, :], os.path.join(folder, f'img_{i_b+i_img}.png'))\n\n        if len(os.listdir('../tmp_images')) != n_images:\n            print(len(os.listdir('../tmp_images')))\n\n        mifid = compute_mifid(folder, DATA_PATH, WEIGHTS_PATH, model_params)\n        shutil.rmtree(folder, ignore_errors=True)\n        return mifid\n\n    \n    def train(self, dataset, epochs, batch_sizes, fade_in_percentage, ema_decays, start_depth=0, verbose=1):\n        assert self.depth == len(batch_sizes), \"batch_sizes not compatible with depth\"\n        infos = {'resolution': [], 'discriminator_loss': [], 'generator_loss': []}\n        self.gen.train()\n        self.dis.train()\n        if self.use_ema:\n            self.gen_shadow.train()\n\n        fixed_noise = torch.randn(128, self.latent_size - self.num_classes).cuda()\n        fixed_races = torch.from_numpy(np.random.choice(range(self.num_classes), size=128)).long()\n\n        for current_depth in range(start_depth, self.depth):\n            current_res = np.power(2, current_depth + 2)\n            print(\"\\n   -> Current resolution: %d x %d \\n\" % (current_res, current_res))\n\n            data = torch.utils.data.DataLoader(dataset, batch_size=batch_sizes[current_depth], num_workers=4, shuffle=True)\n            self.ema_decay = ema_decays[current_depth]\n            ticker = 1\n\n            for epoch in range(1, epochs[current_depth] + 1):\n                start_time = time.time()\n                d_loss = 0\n                g_loss = 0\n                \n                fader_point = fade_in_percentage[current_depth] // 100 * epochs[current_depth] * len(iter(data))\n                step = 0  # counter for number of iterations\n                        \n                if current_res == 64 and (epoch % 50) == 0:\n                    self.ema_decay = 0.9 + self.ema_decay / 10\n\n\n                for (i, batch) in enumerate(data, 1):\n                    # calculate the alpha for fading in the layers\n                    alpha = ticker / fader_point if ticker <= fader_point else 1\n\n                    # extract current batch of data for training\n                    images, labels = batch\n                    images = images.cuda()\n                    labels = labels.view(-1, 1)\n\n                    # create the input to the Generator\n                    label_information = self.one_hot_encode(labels).cuda()\n                    latent_vector = th.randn(images.shape[0], self.latent_size - self.num_classes).cuda()\n                    gan_input = th.cat((label_information, latent_vector), dim=-1)\n\n                    # optimize the discriminator:\n                    dis_loss = self.optimize_discriminator(gan_input, images,\n                                                           labels, current_depth, alpha)\n                    d_loss += dis_loss / len(data)\n                    \n                    # optimize the generator:\n                    gen_loss = self.optimize_generator(gan_input, images,\n                                                       labels, current_depth, alpha)\n                    g_loss += gen_loss / len(data)\n                    \n                    # increment the alpha ticker and the step\n                    ticker += 1\n                    step += 1\n                \n                infos['discriminator_loss'].append(d_loss)\n                infos['generator_loss'].append(g_loss)\n                infos['resolution'].append(current_res)\n                \n                if epoch % verbose == 0:\n                    elapsed_time = time.time() - start_time\n                    print(f'Epoch {epoch}/{epochs[current_depth]}     lr_g={self.lr_g:.1e}     lr_d={self.lr_d:.1e}     ema_decay={self.ema_decay:.4f}', end='     ')\n                    print(f'disc_loss={d_loss:.3f}     gen_loss={g_loss:.3f}     t={elapsed_time:.0f}s')  \n                if epoch % (verbose * 25) == 0 and current_res == 64:\n                    for i in range(5):\n                        self.plot_race(i, depth=current_depth, alpha=alpha, n_plot=5, n=batch_sizes[0])\n#                     score = self.compute_mifid(alpha=alpha)\n#                     print(f'\\n -> MiFID at epoch {epoch} is {score:.3f} \\n')\n                    seed_everything(seed + epoch)\n                elif epoch % (verbose * 10) == 0:\n                    self.generate(current_depth, alpha=alpha, noise=fixed_noise, races=fixed_races, n=batch_sizes[0], n_plot=10)\n                \n                if time.time() - KERNEL_START_TIME > 32000:\n                    print('Time limit reached, interrupting training.')\n                    break\n                    \n        self.gen.eval()\n        self.dis.eval()\n        if self.use_ema:\n            self.gen_shadow.eval()\n        return infos","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"seed_everything(seed)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"depth = 5\nlatent_size = 256","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"loss = Hinge\nlr_d = 6e-3\nlr_g = 6e-3","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pro_gan = ConditionalProGAN(num_classes=nb_classes, depth=depth, latent_size=latent_size, \n                            loss=loss, lr_d=lr_d, lr_g=lr_g,\n                            use_ema=True, use_eql=True, use_spec_norm=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_epochs = [5, 10, 20, 40, 100]\n\nfade_ins = [50, 20, 20, 10, 5]\nbatch_sizes = [64] * 5\nema_decays = [0.9, 0.9, 0.99, 0.99, 0.99]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"infos = pro_gan.train(\n    dataset=dataset,\n    epochs=num_epochs,\n    fade_in_percentage=fade_ins,\n    batch_sizes=batch_sizes,\n    ema_decays=ema_decays,\n    verbose=1\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# save_model_weights(pro_gan.gen, \"gen_weights.pt\")\n# save_model_weights(pro_gan.gen_shadow, \"gen_shadow_weights.pt\")\n# save_model_weights(pro_gan.dis, \"dis_weights.pt\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Loss Curves"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"fig, ax1 = plt.subplots(figsize=(12, 8))\n\nplt.plot(infos['discriminator_loss'], label='discriminator loss', c='darkorange')\nax1.set_ylabel('Discriminator Loss', color='darkorange', size=14)\nax1.tick_params(axis='y', colors='darkorange')\nax1.set_xlabel('Epochs', size=14)\nplt.grid(True)\nplt.legend(loc=(0, 1.01))\n\nax2 = ax1.twinx()\nplt.plot(infos['generator_loss'], label='generator loss', c='dodgerblue')\nax2.set_ylabel('Generator Loss', color='dodgerblue', size=14)\nax2.tick_params(axis='y', colors='dodgerblue')\nplt.legend(loc=(0.84, 1.01))\n\nres = 4\nfor epoch in np.cumsum(num_epochs[:-1]):\n    plt.axvline(epoch, c='r', alpha=0.5)\n    plt.text(x=epoch-10, y=np.max(infos['generator_loss']), s=f'{res}x{res}', bbox=dict(facecolor='red', alpha=0.25))\n    res *= 2\n\nplt.title('Loss evolution', size=15)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Generated Doggos"},{"metadata":{"trusted":true},"cell_type":"code","source":"_ = pro_gan.generate(n_plot=25, n=batch_sizes[0])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Submission"},{"metadata":{},"cell_type":"markdown","source":"## Class sampling probabilities\nI computed the public LB metric generating dogs only of one race to evaluate my GAN on specific races. It turned out that removing the worst races did not improve the score. "},{"metadata":{"trusted":true},"cell_type":"code","source":"use_class_probs = False","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"class_mifids = {0: 59.588, 1: 72.73, 2: 67.551, 3: 55.2313, 4: 106.411, 5: 91.73, 6: 76.835, 7: 116.532, 8: 133.186, 9: 97.613, 10: 60.579, 11: 88.973, 12: 100.450, 13: 70.434, 14: 109.0, 15: 67.517, 16: 70.190, 17: 64.704, 18: 93.84, 19: 85.469, 20: 76.62, 21: 59.845, 22: 88.898, 23: 106.39, 24: 82.955, 25: 65.451, 26: 58.066, 27: 51.016, 28: 86.535, 29: 55.7707, 30: 101.185, 31: 107.090, 32: 72.13, 33: 82.097, 34: 75.549, 35: 101.369, 36: 58.6547, 37: 58.2671, 38: 119.541, 39: 102.502, 40: 50.871, 41: 62.708, 42: 68.374, 43: 91.805, 44: 121.485, 45: 63.633, 46: 81.873, 47: 84.495, 48: 70.738, 49: 86.552, 50: 88.067, 51: 71.89, 52: 94.634, 53: 113.624, 54: 83.190, 55: 98.981, 56: 95.340, 57: 64.076, 58: 117.179, 59: 88.268, 60: 72.986, 61: 98.971, 62: 85.285, 63: 85.808, 64: 79.254, 65: 64.701, 66: 86.365, 67: 70.703, 68: 93.268, 69: 86.699, 70: 52.1906, 71: 113.921, 72: 97.362, 73: 91.260, 74: 86.280, 75: 79.446, 76: 93.833, 77: 90.210, 78: 78.724, 79: 71.74, 80: 63.501, 81: 68.04, 82: 63.95, 83: 64.699, 84: 67.830, 85: 74.525, 86: 95.0, 87: 80.052, 88: 66.01, 89: 82.810, 90: 89.005, 91: 70.746, 92: 71.503, 93: 74.91, 94: 93.378, 95: 82.145, 96: 83.216, 97: 82.569, 98: 65.660, 99: 104.155, 100: 92.105, 101: 81.490, 102: 61.454, 103: 93.271, 104: 75.405, 105: 64.791, 106: 67.661, 107: 96.165, 108: 66.574, 109: 80.197, 110: 89.485, 111: 98.295, 112: 73.864, 113: 59.862, 114: 61.0195, 115: 86.696, 116: 71.2, 117: 91.355, 118: 62.981, 119: 65.444}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"threshold = 200 # All classes are kept\nkept_classes = [k for k in class_mifids.keys() if class_mifids[k] < threshold]\nprint(f'Kept {len(kept_classes)} classes')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"probs = [class_mifids[k] for k in class_mifids.keys()]\nprobs = (probs - np.min(probs)) / np.sum(probs) / 4\nprobs = np.clip((1/120 - probs), 1/1000, 1/50)\nprobs /= np.sum(probs)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Generating images"},{"metadata":{"trusted":true},"cell_type":"code","source":"im_batch_size = 100\nn_images = 10000","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if os.path.exists('../output_images'):\n    shutil.rmtree('../output_images', ignore_errors=True)\nos.mkdir('../output_images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i_batch in tqdm(range(0, n_images, im_batch_size)):\n    if use_class_probs:\n        races = np.random.choice(kept_classes, size=im_batch_size, p=probs)\n        races = torch.from_numpy(races).long()\n        gen_images = pro_gan.generate(races=races, n=im_batch_size)\n    else:\n        gen_images = pro_gan.generate(n=im_batch_size)\n    for i_image in range(gen_images.size(0)):\n        save_image(gen_images[i_image, :, :, :], os.path.join('../output_images', f'img_{i_batch+i_image}.png'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print('Number of generated images :', len(os.listdir('../output_images')))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"shutil.make_archive('images', 'zip', '../output_images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print(f\"Done in {(time.time() - begin)/60:.1f} minutes\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Final words\n\nI switched to ProGAN about 2/3 weeks after the beginning of the competition and quicly managed to hit 32 LB, using a setup I optimized on DCGans and tweaking the learning rate of the original repo. Afterwards I spent the rest of the competition trying to improve this architecture, and only managed to earn 2 points. Probably a bad choice, but as I was not very available during the last few weeks, it was easier for me.\n\nThanks for reading, and thanks to the Kaggle team for the huge work on the competition. \n"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}