{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import datetime\nimport gc\nimport json\nimport math\nimport os\nimport time\nimport shutil\nimport sys\n\nfrom functools import partial\nfrom numba import jit, njit\nfrom PIL import Image\nfrom tqdm import tqdm, trange\n\nimport cv2\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import init\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.nn import Parameter as P\nimport torch.utils.data as data\nfrom torch.utils.data import DataLoader\n\nimport torchvision\nfrom torchvision import datasets\nimport torchvision.transforms as T\n\nimport albumentations as A\nfrom albumentations.pytorch import ToTensor\n\nimport xml.etree.ElementTree as ET","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"INPUT_PATH = '../input/all-dogs'\nANNOT_PATH = '../input/annotation/Annotation'\nOUT_DIR = '../output_images'\n\nIMG_SIZE = 64\nIMG_SIZE_2 = IMG_SIZE * 2\nIMG_CHANNELS = 3\nIMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Layers\n\n# Projection of x onto y\ndef proj(x, y):\n    return torch.mm(y, x.t()) * y / torch.mm(y, y.t())\n\n\n# Orthogonalize x wrt list of vectors ys\ndef gram_schmidt(x, ys):\n    for y in ys:\n        x = x - proj(x, y)\n    return x\n\n\n# Apply num_itrs steps of the power method to estimate top N singular values.\ndef power_iteration(W, u_, update=True, eps=1e-12):\n    # Lists holding singular vectors and values\n    us, vs, svs = [], [], []\n    for i, u in enumerate(u_):\n        # Run one step of the power iteration\n        with torch.no_grad():\n            v = torch.matmul(u, W)\n            # Run Gram-Schmidt to subtract components of all other singular vectors\n            v = F.normalize(gram_schmidt(v, vs), eps=eps)\n            # Add to the list\n            vs += [v]\n            # Update the other singular vector\n            u = torch.matmul(v, W.t())\n            # Run Gram-Schmidt to subtract components of all other singular vectors\n            u = F.normalize(gram_schmidt(u, us), eps=eps)\n            # Add to the list\n            us += [u]\n            if update:\n                u_[i][:] = u\n        # Compute this singular value and add it to the list\n        svs += [torch.squeeze(torch.matmul(torch.matmul(v, W.t()), u.t()))]\n        # svs += [torch.sum(F.linear(u, W.transpose(0, 1)) * v)]\n    return svs, us, vs\n\n\n# Convenience passthrough function\nclass identity(nn.Module):\n    def forward(self, input):\n        return input\n\n\n# Spectral normalization base class\nclass SN(object):\n    def __init__(self, num_svs, num_itrs, num_outputs, transpose=False, eps=1e-12):\n        # Number of power iterations per step\n        self.num_itrs = num_itrs\n        # Number of singular values\n        self.num_svs = num_svs\n        # Transposed?\n        self.transpose = transpose\n        # Epsilon value for avoiding divide-by-0\n        self.eps = eps\n        # Register a singular vector for each sv\n        for i in range(self.num_svs):\n            self.register_buffer('u%d' % i, torch.randn(1, num_outputs))\n            self.register_buffer('sv%d' % i, torch.ones(1))\n\n    # Singular vectors (u side)\n    @property\n    def u(self):\n        return [getattr(self, 'u%d' % i) for i in range(self.num_svs)]\n\n    # Singular values;\n    # note that these buffers are just for logging and are not used in training.\n    @property\n    def sv(self):\n        return [getattr(self, 'sv%d' % i) for i in range(self.num_svs)]\n\n    # Compute the spectrally-normalized weight\n    def W_(self):\n        W_mat = self.weight.view(self.weight.size(0), -1)\n        if self.transpose:\n            W_mat = W_mat.t()\n        # Apply num_itrs power iterations\n        for _ in range(self.num_itrs):\n            svs, us, vs = power_iteration(W_mat, self.u, update=self.training, eps=self.eps)\n            # Update the svs\n        if self.training:\n            with torch.no_grad():  # Make sure to do this in a no_grad() context or you'll get memory leaks!\n                for i, sv in enumerate(svs):\n                    self.sv[i][:] = sv\n        return self.weight / svs[0]\n\n\n# 2D Conv layer with spectral norm\nclass SNConv2d(nn.Conv2d, SN):\n    def __init__(self, in_channels, out_channels, kernel_size, stride=1,\n                 padding=0, dilation=1, groups=1, bias=True,\n                 num_svs=1, num_itrs=1, eps=1e-12):\n        nn.Conv2d.__init__(self, in_channels, out_channels, kernel_size, stride,\n                           padding, dilation, groups, bias)\n        SN.__init__(self, num_svs, num_itrs, out_channels, eps=eps)\n\n    def forward(self, x):\n        return F.conv2d(x, self.W_(), self.bias, self.stride,\n                        self.padding, self.dilation, self.groups)\n\n\n# Linear layer with spectral norm\nclass SNLinear(nn.Linear, SN):\n    def __init__(self, in_features, out_features, bias=True,\n                 num_svs=1, num_itrs=1, eps=1e-12):\n        nn.Linear.__init__(self, in_features, out_features, bias)\n        SN.__init__(self, num_svs, num_itrs, out_features, eps=eps)\n\n    def forward(self, x):\n        return F.linear(x, self.W_(), self.bias)\n\n\n# Embedding layer with spectral norm\n# We use num_embeddings as the dim instead of embedding_dim here\n# for convenience sake\nclass SNEmbedding(nn.Embedding, SN):\n    def __init__(self, num_embeddings, embedding_dim, padding_idx=None,\n                 max_norm=None, norm_type=2, scale_grad_by_freq=False,\n                 sparse=False, _weight=None,\n                 num_svs=1, num_itrs=1, eps=1e-12):\n        nn.Embedding.__init__(self, num_embeddings, embedding_dim, padding_idx,\n                              max_norm, norm_type, scale_grad_by_freq,\n                              sparse, _weight)\n        SN.__init__(self, num_svs, num_itrs, num_embeddings, eps=eps)\n\n    def forward(self, x):\n        return F.embedding(x, self.W_())\n\n\n# A non-local block as used in SA-GAN\n# Note that the implementation as described in the paper is largely incorrect;\n# refer to the released code for the actual implementation.\n# class Attention(nn.Module):\n#     def __init__(self, ch, which_conv=SNConv2d, name='attention'):\n#         super(Attention, self).__init__()\n#         # Channel multiplier\n#         self.ch = ch\n#         self.which_conv = which_conv\n#         self.theta = self.which_conv(self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False)\n#         self.phi = self.which_conv(self.ch, self.ch // 8, kernel_size=1, padding=0, bias=False)\n#         self.g = self.which_conv(self.ch, self.ch // 2, kernel_size=1, padding=0, bias=False)\n#         self.o = self.which_conv(self.ch // 2, self.ch, kernel_size=1, padding=0, bias=False)\n#         # Learnable gain parameter\n#         self.gamma = P(torch.tensor(0.), requires_grad=True)\n\n#     def forward(self, x, y=None):\n#         # Apply convs\n#         theta = self.theta(x)\n#         phi = F.max_pool2d(self.phi(x), [2, 2])\n#         g = F.max_pool2d(self.g(x), [2, 2])\n#         # Perform reshapes\n#         theta = theta.view(-1, self.ch // 8, x.shape[2] * x.shape[3])\n#         phi = phi.view(-1, self.ch // 8, x.shape[2] * x.shape[3] // 4)\n#         g = g.view(-1, self.ch // 2, x.shape[2] * x.shape[3] // 4)\n#         # Matmul and softmax to get attention maps\n#         beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), -1)\n#         # Attention map times g path\n#         o = self.o(torch.bmm(g, beta.transpose(1, 2)).view(-1, self.ch // 2, x.shape[2], x.shape[3]))\n#         return self.gamma * o + x\n    \n    \nclass Attention(nn.Module):\n    def __init__(self, ch, which_conv=SNConv2d, name='attention'):\n        super(Attention, self).__init__()\n        # Channel multiplier\n        self.ch = ch\n        self.which_conv = which_conv\n        self.theta = self.which_conv(self.ch, self.ch // 4, kernel_size=1, padding=0, bias=False)\n        self.phi = self.which_conv(self.ch, self.ch // 4, kernel_size=1, padding=0, bias=False)\n        self.g = self.which_conv(self.ch, self.ch // 2, kernel_size=1, padding=0, bias=False)\n        self.o = self.which_conv(self.ch // 2, self.ch, kernel_size=1, padding=0, bias=False)\n        # Learnable gain parameter\n        self.gamma = P(torch.tensor(0.), requires_grad=True)\n\n    def forward(self, x, y=None):\n        # Apply convs\n        theta = self.theta(x)\n        phi = F.max_pool2d(self.phi(x), [2, 2])\n        g = F.max_pool2d(self.g(x), [2, 2])\n        # Perform reshapes\n        theta = theta.view(-1, self.ch // 4, x.shape[2] * x.shape[3])\n        phi = phi.view(-1, self.ch // 4, x.shape[2] * x.shape[3] // 4)\n        g = g.view(-1, self.ch // 2, x.shape[2] * x.shape[3] // 4)\n        # Matmul and softmax to get attention maps\n        beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), -1)\n        # Attention map times g path\n        o = self.o(torch.bmm(g, beta.transpose(1, 2)).view(-1, self.ch // 2, x.shape[2], x.shape[3]))\n        return self.gamma * o + x\n\n\n# Fused batchnorm op\ndef fused_bn(x, mean, var, gain=None, bias=None, eps=1e-5):\n    # Apply scale and shift--if gain and bias are provided, fuse them here\n    # Prepare scale\n    scale = torch.rsqrt(var + eps)\n    # If a gain is provided, use it\n    if gain is not None:\n        scale = scale * gain\n    # Prepare shift\n    shift = mean * scale\n    # If bias is provided, use it\n    if bias is not None:\n        shift = shift - bias\n    return x * scale - shift\n    # return ((x - mean) / ((var + eps) ** 0.5)) * gain + bias # The unfused way.\n\n\n# Manual BN\n# Calculate means and variances using mean-of-squares minus mean-squared\ndef manual_bn(x, gain=None, bias=None, return_mean_var=False, eps=1e-5):\n    # Cast x to float32 if necessary\n    float_x = x.float()\n    # Calculate expected value of x (m) and expected value of x**2 (m2)\n    # Mean of x\n    m = torch.mean(float_x, [0, 2, 3], keepdim=True)\n    # Mean of x squared\n    m2 = torch.mean(float_x ** 2, [0, 2, 3], keepdim=True)\n    # Calculate variance as mean of squared minus mean squared.\n    var = (m2 - m ** 2)\n    # Cast back to float 16 if necessary\n    var = var.type(x.type())\n    m = m.type(x.type())\n    # Return mean and variance for updating stored mean/var if requested\n    if return_mean_var:\n        return fused_bn(x, m, var, gain, bias, eps), m.squeeze(), var.squeeze()\n    else:\n        return fused_bn(x, m, var, gain, bias, eps)\n\n\n# My batchnorm, supports standing stats\nclass myBN(nn.Module):\n    def __init__(self, num_channels, eps=1e-5, momentum=0.1):\n        super(myBN, self).__init__()\n        # momentum for updating running stats\n        self.momentum = momentum\n        # epsilon to avoid dividing by 0\n        self.eps = eps\n        # Momentum\n        self.momentum = momentum\n        # Register buffers\n        self.register_buffer('stored_mean', torch.zeros(num_channels))\n        self.register_buffer('stored_var', torch.ones(num_channels))\n        self.register_buffer('accumulation_counter', torch.zeros(1))\n        # Accumulate running means and vars\n        self.accumulate_standing = False\n\n    # reset standing stats\n    def reset_stats(self):\n        self.stored_mean[:] = 0\n        self.stored_var[:] = 0\n        self.accumulation_counter[:] = 0\n\n    def forward(self, x, gain, bias):\n        if self.training:\n            out, mean, var = manual_bn(x, gain, bias, return_mean_var=True, eps=self.eps)\n            # If accumulating standing stats, increment them\n            if self.accumulate_standing:\n                self.stored_mean[:] = self.stored_mean + mean.data\n                self.stored_var[:] = self.stored_var + var.data\n                self.accumulation_counter += 1.0\n            # If not accumulating standing stats, take running averages\n            else:\n                self.stored_mean[:] = self.stored_mean * (1 - self.momentum) + mean * self.momentum\n                self.stored_var[:] = self.stored_var * (1 - self.momentum) + var * self.momentum\n            return out\n        # If not in training mode, use the stored statistics\n        else:\n            mean = self.stored_mean.view(1, -1, 1, 1)\n            var = self.stored_var.view(1, -1, 1, 1)\n            # If using standing stats, divide them by the accumulation counter\n            if self.accumulate_standing:\n                mean = mean / self.accumulation_counter\n                var = var / self.accumulation_counter\n            return fused_bn(x, mean, var, gain, bias, self.eps)\n\n\n# Simple function to handle groupnorm norm stylization\ndef groupnorm(x, norm_style):\n    # If number of channels specified in norm_style:\n    if 'ch' in norm_style:\n        ch = int(norm_style.split('_')[-1])\n        groups = max(int(x.shape[1]) // ch, 1)\n    # If number of groups specified in norm style\n    elif 'grp' in norm_style:\n        groups = int(norm_style.split('_')[-1])\n    # If neither, default to groups = 16\n    else:\n        groups = 16\n    return F.group_norm(x, groups)\n\n\n# Class-conditional bn\n# output size is the number of channels, input size is for the linear layers\n# Andy's Note: this class feels messy but I'm not really sure how to clean it up\n# Suggestions welcome! (By which I mean, refactor this and make a pull request\n# if you want to make this more readable/usable).\nclass ccbn(nn.Module):\n    def __init__(self, output_size, input_size, which_linear, eps=1e-5, momentum=0.1,\n                 cross_replica=False, mybn=False, norm_style='bn', ):\n        super(ccbn, self).__init__()\n        self.output_size, self.input_size = output_size, input_size\n        # Prepare gain and bias layers\n        self.gain = which_linear(input_size, output_size)\n        self.bias = which_linear(input_size, output_size)\n        # epsilon to avoid dividing by 0\n        self.eps = eps\n        # Momentum\n        self.momentum = momentum\n        # Use cross-replica batchnorm?\n        self.cross_replica = cross_replica\n        # Use my batchnorm?\n        self.mybn = mybn\n        # Norm style?\n        self.norm_style = norm_style\n\n        if self.mybn:\n            self.bn = myBN(output_size, self.eps, self.momentum)\n        elif self.norm_style in ['bn', 'in']:\n            self.register_buffer('stored_mean', torch.zeros(output_size))\n            self.register_buffer('stored_var', torch.ones(output_size))\n\n    def forward(self, x, y):\n        # Calculate class-conditional gains and biases\n        gain = (1 + self.gain(y)).view(y.size(0), -1, 1, 1)\n        bias = self.bias(y).view(y.size(0), -1, 1, 1)\n        # If using my batchnorm\n        if self.mybn:\n            return self.bn(x, gain=gain, bias=bias)\n        # else:\n        else:\n            if self.norm_style == 'bn':\n                out = F.batch_norm(x, self.stored_mean, self.stored_var, None, None,\n                                   self.training, 0.1, self.eps)\n            elif self.norm_style == 'in':\n                out = F.instance_norm(x, self.stored_mean, self.stored_var, None, None,\n                                      self.training, 0.1, self.eps)\n            elif self.norm_style == 'gn':\n                out = groupnorm(x, self.normstyle)\n            elif self.norm_style == 'nonorm':\n                out = x\n            return out * gain + bias\n\n    def extra_repr(self):\n        s = 'out: {output_size}, in: {input_size},'\n        s += ' cross_replica={cross_replica}'\n        return s.format(**self.__dict__)\n\n\n# Normal, non-class-conditional BN\nclass bn(nn.Module):\n    def __init__(self, output_size, eps=1e-5, momentum=0.1,\n                 cross_replica=False, mybn=False):\n        super(bn, self).__init__()\n        self.output_size = output_size\n        # Prepare gain and bias layers\n        self.gain = P(torch.ones(output_size), requires_grad=True)\n        self.bias = P(torch.zeros(output_size), requires_grad=True)\n        # epsilon to avoid dividing by 0\n        self.eps = eps\n        # Momentum\n        self.momentum = momentum\n        # Use cross-replica batchnorm?\n        self.cross_replica = cross_replica\n        # Use my batchnorm?\n        self.mybn = mybn\n\n        if mybn:\n            self.bn = myBN(output_size, self.eps, self.momentum)\n        # Register buffers if neither of the above\n        else:\n            self.register_buffer('stored_mean', torch.zeros(output_size))\n            self.register_buffer('stored_var', torch.ones(output_size))\n\n    def forward(self, x, y=None):\n        if self.mybn:\n            gain = self.gain.view(1, -1, 1, 1)\n            bias = self.bias.view(1, -1, 1, 1)\n            return self.bn(x, gain=gain, bias=bias)\n        else:\n            return F.batch_norm(x, self.stored_mean, self.stored_var, self.gain,\n                                self.bias, self.training, self.momentum, self.eps)\n\n\n# Generator blocks\n# Note that this class assumes the kernel size and padding (and any other\n# settings) have been selected in the main generator module and passed in\n# through the which_conv arg. Similar rules apply with which_bn (the input\n# size [which is actually the number of channels of the conditional info] must\n# be preselected)\nclass GBlock(nn.Module):\n    def __init__(self, in_channels, out_channels,\n                 which_conv=nn.Conv2d, which_bn=bn, activation=None,\n                 upsample=None):\n        super(GBlock, self).__init__()\n\n        self.in_channels, self.out_channels = in_channels, out_channels\n        self.which_conv, self.which_bn = which_conv, which_bn\n        self.activation = activation\n        self.upsample = upsample\n        # Conv layers\n        self.conv1 = self.which_conv(self.in_channels, self.out_channels)\n        self.conv2 = self.which_conv(self.out_channels, self.out_channels)\n        self.learnable_sc = in_channels != out_channels or upsample\n        if self.learnable_sc:\n            self.conv_sc = self.which_conv(in_channels, out_channels,\n                                           kernel_size=1, padding=0)\n        # Batchnorm layers\n        self.bn1 = self.which_bn(in_channels)\n        self.bn2 = self.which_bn(out_channels)\n        # upsample layers\n        self.upsample = upsample\n\n    def forward(self, x, y):\n        h = self.activation(self.bn1(x, y))\n        if self.upsample:\n            h = self.upsample(h)\n            x = self.upsample(x)\n        h = self.conv1(h)\n        h = self.activation(self.bn2(h, y))\n        h = self.conv2(h)\n        if self.learnable_sc:\n            x = self.conv_sc(x)\n        return h + x\n\n\n# Residual block for the discriminator\nclass DBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, which_conv=SNConv2d, wide=True,\n                 preactivation=False, activation=None, downsample=None, ):\n        super(DBlock, self).__init__()\n        self.in_channels, self.out_channels = in_channels, out_channels\n        # If using wide D (as in SA-GAN and BigGAN), change the channel pattern\n        self.hidden_channels = self.out_channels if wide else self.in_channels\n        self.which_conv = which_conv\n        self.preactivation = preactivation\n        self.activation = activation\n        self.downsample = downsample\n\n        # Conv layers\n        self.conv1 = self.which_conv(self.in_channels, self.hidden_channels)\n        self.conv2 = self.which_conv(self.hidden_channels, self.out_channels)\n        self.learnable_sc = True if (in_channels != out_channels) or downsample else False\n        if self.learnable_sc:\n            self.conv_sc = self.which_conv(in_channels, out_channels,\n                                           kernel_size=1, padding=0)\n\n    def shortcut(self, x):\n        if self.preactivation:\n            if self.learnable_sc:\n                x = self.conv_sc(x)\n            if self.downsample:\n                x = self.downsample(x)\n        else:\n            if self.downsample:\n                x = self.downsample(x)\n            if self.learnable_sc:\n                x = self.conv_sc(x)\n        return x\n\n    def forward(self, x):\n        if self.preactivation:\n            # h = self.activation(x) # NOT TODAY SATAN\n            # Andy's note: This line *must* be an out-of-place ReLU or it\n            #              will negatively affect the shortcut connection.\n            h = F.relu(x)\n        else:\n            h = x\n        h = self.conv1(h)\n        h = self.conv2(self.activation(h))\n        if self.downsample:\n            h = self.downsample(h)\n\n        return h + self.shortcut(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Architectures for G\n# Attention is passed in in the format '32_64' to mean applying an attention\n# block at both resolution 32x32 and 64x64. Just '64' will apply at 64x64.\ndef G_arch(ch=64, attention='64', ksize='333333', dilation='111111'):\n    arch = {}\n    arch[512] = {'in_channels': [ch * item for item in [16, 16, 8, 8, 4, 2, 1]],\n                 'out_channels': [ch * item for item in [16, 8, 8, 4, 2, 1, 1]],\n                 'upsample': [True] * 7,\n                 'resolution': [8, 16, 32, 64, 128, 256, 512],\n                 'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])\n                               for i in range(3, 10)}}\n    arch[256] = {'in_channels': [ch * item for item in [16, 16, 8, 8, 4, 2]],\n                 'out_channels': [ch * item for item in [16, 8, 8, 4, 2, 1]],\n                 'upsample': [True] * 6,\n                 'resolution': [8, 16, 32, 64, 128, 256],\n                 'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])\n                               for i in range(3, 9)}}\n    arch[128] = {'in_channels': [ch * item for item in [16, 16, 8, 4, 2]],\n                 'out_channels': [ch * item for item in [16, 8, 4, 2, 1]],\n                 'upsample': [True] * 5,\n                 'resolution': [8, 16, 32, 64, 128],\n                 'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])\n                               for i in range(3, 8)}}\n    arch[64] = {'in_channels': [ch * item for item in [16, 16, 8, 4]],\n                'out_channels': [ch * item for item in [16, 8, 4, 2]],\n                'upsample': [True] * 4,\n                'resolution': [8, 16, 32, 64],\n                'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])\n                              for i in range(3, 7)}}\n    arch[32] = {'in_channels': [ch * item for item in [4, 4, 4]],\n                'out_channels': [ch * item for item in [4, 4, 4]],\n                'upsample': [True] * 3,\n                'resolution': [8, 16, 32],\n                'attention': {2 ** i: (2 ** i in [int(item) for item in attention.split('_')])\n                              for i in range(3, 6)}}\n\n    return arch\n\n\nclass Generator(nn.Module):\n    def __init__(self, G_ch=64, dim_z=128, bottom_width=4, resolution=128,\n                 G_kernel_size=3, G_attn='64', n_classes=1000,\n                 num_G_SVs=1, num_G_SV_itrs=1,\n                 G_shared=True, shared_dim=0, hier=False,\n                 cross_replica=False, mybn=False,\n                 G_activation=nn.ReLU(inplace=False),\n                 G_lr=5e-5, G_B1=0.0, G_B2=0.999, adam_eps=1e-8,\n                 BN_eps=1e-5, SN_eps=1e-12, G_mixed_precision=False, G_fp16=False,\n                 G_init='ortho', skip_init=False, no_optim=False,\n                 G_param='SN', norm_style='bn',\n                 **kwargs):\n        super(Generator, self).__init__()\n        # Channel width mulitplier\n        self.ch = G_ch\n        # Dimensionality of the latent space\n        self.dim_z = dim_z\n        # The initial spatial dimensions\n        self.bottom_width = bottom_width\n        # Resolution of the output\n        self.resolution = resolution\n        # Kernel size?\n        self.kernel_size = G_kernel_size\n        # Attention?\n        self.attention = G_attn\n        # number of classes, for use in categorical conditional generation\n        self.n_classes = n_classes\n        # Use shared embeddings?\n        self.G_shared = G_shared\n        # Dimensionality of the shared embedding? Unused if not using G_shared\n        self.shared_dim = shared_dim if shared_dim > 0 else dim_z\n        # Hierarchical latent space?\n        self.hier = hier\n        # Cross replica batchnorm?\n        self.cross_replica = cross_replica\n        # Use my batchnorm?\n        self.mybn = mybn\n        # nonlinearity for residual blocks\n        self.activation = G_activation\n        # Initialization style\n        self.init = G_init\n        # Parameterization style\n        self.G_param = G_param\n        # Normalization style\n        self.norm_style = norm_style\n        # Epsilon for BatchNorm?\n        self.BN_eps = BN_eps\n        # Epsilon for Spectral Norm?\n        self.SN_eps = SN_eps\n        # fp16?\n        self.fp16 = G_fp16\n        # Architecture dict\n        self.arch = G_arch(self.ch, self.attention)[resolution]\n\n        # If using hierarchical latents, adjust z\n        if self.hier:\n            # Number of places z slots into\n            self.num_slots = len(self.arch['in_channels']) + 1\n            self.z_chunk_size = (self.dim_z // self.num_slots)\n            # Recalculate latent dimensionality for even splitting into chunks\n            self.dim_z = self.z_chunk_size * self.num_slots\n        else:\n            self.num_slots = 1\n            self.z_chunk_size = 0\n\n        # Which convs, batchnorms, and linear layers to use\n        if self.G_param == 'SN':\n            self.which_conv = partial(SNConv2d,\n                kernel_size=3, padding=1,\n                num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,\n                eps=self.SN_eps)\n            self.which_linear = partial(SNLinear,\n                num_svs=num_G_SVs, num_itrs=num_G_SV_itrs,\n                eps=self.SN_eps)\n        else:\n            self.which_conv = partial(nn.Conv2d, kernel_size=3, padding=1)\n            self.which_linear = nn.Linear\n\n        # We use a non-spectral-normed embedding here regardless;\n        # For some reason applying SN to G's embedding seems to randomly cripple G\n        self.which_embedding = nn.Embedding\n        bn_linear = (partial(self.which_linear, bias=False) if self.G_shared else self.which_embedding)\n        self.which_bn = partial(ccbn,\n            which_linear=bn_linear,\n            cross_replica=self.cross_replica,\n            mybn=self.mybn,\n            input_size=(self.shared_dim + self.z_chunk_size if self.G_shared else self.n_classes),\n            norm_style=self.norm_style,\n            eps=self.BN_eps)\n\n        # Prepare model\n        # If not using shared embeddings, self.shared is just a passthrough\n        self.shared = (self.which_embedding(n_classes, self.shared_dim) if G_shared else identity())\n        # First linear layer\n        self.linear = self.which_linear(\n            self.dim_z // self.num_slots,\n            self.arch['in_channels'][0] * (self.bottom_width ** 2))\n\n        # self.blocks is a doubly-nested list of modules, the outer loop intended\n        # to be over blocks at a given resolution (resblocks and/or self-attention)\n        # while the inner loop is over a given block\n        self.blocks = []\n        for index in range(len(self.arch['out_channels'])):\n            self.blocks += [[GBlock(\n                in_channels=self.arch['in_channels'][index],\n                out_channels=self.arch['out_channels'][index],\n                which_conv=self.which_conv,\n                which_bn=self.which_bn,\n                activation=self.activation,\n                upsample=(partial(F.interpolate, scale_factor=2) if self.arch['upsample'][index] else None)\n            )]]\n\n            # If attention on this block, attach it to the end\n            if self.arch['attention'][self.arch['resolution'][index]]:\n                print('Adding attention layer in G at resolution %d' % self.arch['resolution'][index])\n                self.blocks[-1] += [Attention(self.arch['out_channels'][index], self.which_conv)]\n\n        # Turn self.blocks into a ModuleList so that it's all properly registered.\n        self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])\n\n        # output layer: batchnorm-relu-conv.\n        # Consider using a non-spectral conv here\n        self.output_layer = nn.Sequential(\n            bn(self.arch['out_channels'][-1], cross_replica=self.cross_replica, mybn=self.mybn),\n            self.activation,\n            self.which_conv(self.arch['out_channels'][-1], 3)\n        )\n\n        # Initialize weights. Optionally skip init for testing.\n        if not skip_init:\n            self.init_weights()\n\n        # Set up optimizer\n        # If this is an EMA copy, no need for an optim, so just return now\n        if no_optim:\n            return\n        self.lr, self.B1, self.B2, self.adam_eps = G_lr, G_B1, G_B2, adam_eps\n        if G_mixed_precision:\n            print('Using fp16 adam in G...')\n            self.optim = Adam16(params=self.parameters(), lr=self.lr,\n                                betas=(self.B1, self.B2), weight_decay=0,\n                                eps=self.adam_eps)\n        else:\n            self.optim = optim.Adam(params=self.parameters(), lr=self.lr,\n                                    betas=(self.B1, self.B2), weight_decay=0,\n                                    eps=self.adam_eps)\n\n        # LR scheduling, left here for forward compatibility\n        # self.lr_sched = {'itr' : 0}# if self.progressive else {}\n        # self.j = 0\n\n    # Initialize\n    def init_weights(self):\n        self.param_count = 0\n        for module in self.modules():\n            if (isinstance(module, nn.Conv2d)\n                    or isinstance(module, nn.Linear)\n                    or isinstance(module, nn.Embedding)):\n                if self.init == 'ortho':\n                    init.orthogonal_(module.weight)\n                elif self.init == 'N02':\n                    init.normal_(module.weight, 0, 0.02)\n                elif self.init in ['glorot', 'xavier']:\n                    init.xavier_uniform_(module.weight)\n                else:\n                    print('Init style not recognized...')\n                self.param_count += sum([p.data.nelement() for p in module.parameters()])\n        print('Param count for G''s initialized parameters: %d' % self.param_count)\n\n    # Note on this forward function: we pass in a y vector which has\n    # already been passed through G.shared to enable easy class-wise\n    # interpolation later. If we passed in the one-hot and then ran it through\n    # G.shared in this forward function, it would be harder to handle.\n    def forward(self, z, y):\n        # If hierarchical, concatenate zs and ys\n        if self.hier:\n            zs = torch.split(z, self.z_chunk_size, 1)\n            z = zs[0]\n            ys = [torch.cat([y, item], 1) for item in zs[1:]]\n        else:\n            ys = [y] * len(self.blocks)\n\n        # First linear layer\n        h = self.linear(z)\n        # Reshape\n        h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width)\n\n        # Loop over blocks\n        for index, blocklist in enumerate(self.blocks):\n            # Second inner loop in case block has multiple layers\n            for block in blocklist:\n                h = block(h, ys[index])\n\n        # Apply batchnorm-relu-conv-tanh at output\n        return torch.tanh(self.output_layer(h))\n\n\n# Discriminator architecture, same paradigm as G's above\ndef D_arch(ch=64, attention='64', ksize='333333', dilation='111111'):\n    arch = {}\n    arch[256] = {'in_channels': [3] + [ch * item for item in [1, 2, 4, 8, 8, 16]],\n                 'out_channels': [item * ch for item in [1, 2, 4, 8, 8, 16, 16]],\n                 'downsample': [True] * 6 + [False],\n                 'resolution': [128, 64, 32, 16, 8, 4, 4],\n                 'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]\n                               for i in range(2, 8)}}\n    arch[128] = {'in_channels': [3] + [ch * item for item in [1, 2, 4, 8, 16]],\n                 'out_channels': [item * ch for item in [1, 2, 4, 8, 16, 16]],\n                 'downsample': [True] * 5 + [False],\n                 'resolution': [64, 32, 16, 8, 4, 4],\n                 'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]\n                               for i in range(2, 8)}}\n    arch[64] = {'in_channels': [3] + [ch * item for item in [1, 2, 4, 8]],\n                'out_channels': [item * ch for item in [1, 2, 4, 8, 16]],\n                'downsample': [True] * 4 + [False],\n                'resolution': [32, 16, 8, 4, 4],\n                'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]\n                              for i in range(2, 7)}}\n    arch[32] = {'in_channels': [3] + [item * ch for item in [4, 4, 4]],\n                'out_channels': [item * ch for item in [4, 4, 4, 4]],\n                'downsample': [True, True, False, False],\n                'resolution': [16, 16, 16, 16],\n                'attention': {2 ** i: 2 ** i in [int(item) for item in attention.split('_')]\n                              for i in range(2, 6)}}\n    return arch\n\n\nclass Discriminator(nn.Module):\n\n    def __init__(self, D_ch=64, D_wide=True, resolution=128,\n                 D_kernel_size=3, D_attn='64', n_classes=1000,\n                 num_D_SVs=1, num_D_SV_itrs=1, D_activation=nn.ReLU(inplace=False),\n                 D_lr=2e-4, D_B1=0.0, D_B2=0.999, adam_eps=1e-8,\n                 SN_eps=1e-12, output_dim=1, D_mixed_precision=False, D_fp16=False,\n                 D_init='ortho', skip_init=False, D_param='SN', **kwargs):\n        super(Discriminator, self).__init__()\n        # Width multiplier\n        self.ch = D_ch\n        # Use Wide D as in BigGAN and SA-GAN or skinny D as in SN-GAN?\n        self.D_wide = D_wide\n        # Resolution\n        self.resolution = resolution\n        # Kernel size\n        self.kernel_size = D_kernel_size\n        # Attention?\n        self.attention = D_attn\n        # Number of classes\n        self.n_classes = n_classes\n        # Activation\n        self.activation = D_activation\n        # Initialization style\n        self.init = D_init\n        # Parameterization style\n        self.D_param = D_param\n        # Epsilon for Spectral Norm?\n        self.SN_eps = SN_eps\n        # Fp16?\n        self.fp16 = D_fp16\n        # Architecture\n        self.arch = D_arch(self.ch, self.attention)[resolution]\n\n        # Which convs, batchnorms, and linear layers to use\n        # No option to turn off SN in D right now\n        if self.D_param == 'SN':\n            self.which_conv = partial(SNConv2d,\n                kernel_size=3, padding=1,\n                num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,\n                eps=self.SN_eps)\n            self.which_linear = partial(SNLinear,\n                num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,\n                eps=self.SN_eps)\n            self.which_embedding = partial(SNEmbedding,\n                num_svs=num_D_SVs, num_itrs=num_D_SV_itrs,\n                eps=self.SN_eps)\n        # Prepare model\n        # self.blocks is a doubly-nested list of modules, the outer loop intended\n        # to be over blocks at a given resolution (resblocks and/or self-attention)\n        self.blocks = []\n        for index in range(len(self.arch['out_channels'])):\n            self.blocks += [[DBlock(\n                in_channels=self.arch['in_channels'][index],\n                out_channels=self.arch['out_channels'][index],\n                which_conv=self.which_conv,\n                wide=self.D_wide,\n                activation=self.activation,\n                preactivation=(index > 0),\n                downsample=(nn.AvgPool2d(2) if self.arch['downsample'][index] else None)\n            )]]\n            # If attention on this block, attach it to the end\n            if self.arch['attention'][self.arch['resolution'][index]]:\n                print('Adding attention layer in D at resolution %d' % self.arch['resolution'][index])\n                self.blocks[-1] += [Attention(self.arch['out_channels'][index], self.which_conv)]\n        # Turn self.blocks into a ModuleList so that it's all properly registered.\n        self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])\n        # Linear output layer. The output dimension is typically 1, but may be\n        # larger if we're e.g. turning this into a VAE with an inference output\n        self.linear = self.which_linear(self.arch['out_channels'][-1], output_dim)\n        # Embedding for projection discrimination\n        self.embed = self.which_embedding(self.n_classes, self.arch['out_channels'][-1])\n\n        # Initialize weights\n        if not skip_init:\n            self.init_weights()\n\n        # Set up optimizer\n        self.lr, self.B1, self.B2, self.adam_eps = D_lr, D_B1, D_B2, adam_eps\n        if D_mixed_precision:\n            print('Using fp16 adam in D...')\n            self.optim = Adam16(params=self.parameters(), lr=self.lr,\n                                betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)\n        else:\n            self.optim = optim.Adam(params=self.parameters(), lr=self.lr,\n                                    betas=(self.B1, self.B2), weight_decay=0, eps=self.adam_eps)\n        # LR scheduling, left here for forward compatibility\n        # self.lr_sched = {'itr' : 0}# if self.progressive else {}\n        # self.j = 0\n\n    # Initialize\n    def init_weights(self):\n        self.param_count = 0\n        for module in self.modules():\n            if (isinstance(module, nn.Conv2d)\n                    or isinstance(module, nn.Linear)\n                    or isinstance(module, nn.Embedding)):\n                if self.init == 'ortho':\n                    init.orthogonal_(module.weight)\n                elif self.init == 'N02':\n                    init.normal_(module.weight, 0, 0.02)\n                elif self.init in ['glorot', 'xavier']:\n                    init.xavier_uniform_(module.weight)\n                else:\n                    print('Init style not recognized...')\n                self.param_count += sum([p.data.nelement() for p in module.parameters()])\n        print('Param count for D''s initialized parameters: %d' % self.param_count)\n\n    def forward(self, x, y=None):\n        # Stick x into h for cleaner for loops without flow control\n        h = x\n        # Loop over blocks\n        for index, blocklist in enumerate(self.blocks):\n            for block in blocklist:\n                h = block(h)\n        # Apply global sum pooling as in SN-GAN\n        h = torch.sum(self.activation(h), [2, 3])\n        # Get initial class-unconditional output\n        out = self.linear(h)\n        # Get projection of final featureset onto class vectors and add to evidence\n        out = out + torch.sum(self.embed(y) * h, 1, keepdim=True)\n        return out\n\n\n# Parallelized G_D to minimize cross-gpu communication\n# Without this, Generator outputs would get all-gathered and then rebroadcast.\nclass G_D(nn.Module):\n    def __init__(self, G, D):\n        super(G_D, self).__init__()\n        self.G = G\n        self.D = D\n\n    def forward(self, z, gy, x=None, dy=None, train_G=False, return_G_z=False,\n                split_D=False):\n        # If training G, enable grad tape\n        with torch.set_grad_enabled(train_G):\n            # Get Generator output given noise\n            G_z = self.G(z, self.G.shared(gy))\n            # Cast as necessary\n            if self.G.fp16 and not self.D.fp16:\n                G_z = G_z.float()\n            if self.D.fp16 and not self.G.fp16:\n                G_z = G_z.half()\n        # Split_D means to run D once with real data and once with fake,\n        # rather than concatenating along the batch dimension.\n        if split_D:\n            D_fake = self.D(G_z, gy)\n            if x is not None:\n                D_real = self.D(x, dy)\n                return D_fake, D_real\n            else:\n                if return_G_z:\n                    return D_fake, G_z\n                else:\n                    return D_fake\n        # If real data is provided, concatenate it with the Generator's output\n        # along the batch dimension for improved efficiency.\n        else:\n            D_input = torch.cat([G_z, x], 0) if x is not None else G_z\n            D_class = torch.cat([gy, dy], 0) if dy is not None else gy\n            # Get Discriminator output\n            D_out = self.D(D_input, D_class)\n            if x is not None:\n                return torch.split(D_out, [G_z.shape[0], x.shape[0]])  # D_fake, D_real\n            else:\n                if return_G_z:\n                    return D_out, G_z\n                else:\n                    return D_out","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class MyImg():\n    def __init__(self, img, tfm):\n        self.px = np.array(img)\n        self.tfm = tfm\n        \n    @property\n    def size(self):\n        h, w, _ = self.px.shape\n        return min(w, h)\n\n\ndef pad(img, padding_mode='reflect'):\n    p = math.ceil((max(img.size) - min(img.size)) / 2)\n    p_horr = p if img.width < img.height else 0\n    p_vert = p if img.height < img.width else 0\n    img = T.Pad((p_horr, p_vert), padding_mode=padding_mode)(img)\n    if img.width != img.height:\n        s = min(img.size)\n        img = img.crop((0, 0, s, s))\n    return img\n\n\ndef take_top(img):\n    size = min(img.size)\n    bbox = (0, 0, size, size)\n    return img.crop(bbox)\n\n\ndef take_diagonal(img):\n    w, h = img.size\n    size = min(w, h)\n    bbox_l = (0, 0, size, size)\n    bbox_r = (w - size, h - size, w, h)\n    return [img.crop(bbox_l), img.crop(bbox_r)]\n    \n\nresize = T.Resize(IMG_SIZE, interpolation=Image.LANCZOS)\nresize2x = T.Resize(IMG_SIZE_2, interpolation=Image.LANCZOS)\n\ncenter_crop = T.Compose([resize, T.CenterCrop(IMG_SIZE)])\ncenter_crop2x = T.Compose([resize2x, T.CenterCrop(IMG_SIZE_2)])\n\ntop_crop = T.Compose([T.Lambda(take_top), resize])\ntop_crop2x = T.Compose([T.Lambda(take_top), resize2x])\n\ntwo_crops = T.Compose([\n    resize, T.FiveCrop(IMG_SIZE),\n    T.Lambda(lambda imgs: [imgs[i] for i in [0, 2]])\n])\ntwo_crops2x = T.Compose([resize2x, T.Lambda(take_diagonal)])\npad_only = T.Compose([T.Lambda(pad), resize])\npad_only2x = T.Compose([T.Lambda(pad), resize2x])\n\n\n@jit(nopython=True)\ndef pad_one_dim(clow, chigh, pad, cmax):\n    clow = max(0, clow - pad)\n    chigh = min(cmax, chigh + pad)\n    return clow, chigh, chigh - clow\n\n\ndef calc_bbox(obj, img_w, img_h, zoom=0.0, try_square=True):\n    bndbox = obj.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    \n    # occasionally i get bboxes which exceed img size\n    xmin, xmax, obj_w = pad_one_dim(xmin, xmax, 0, img_w)\n    ymin, ymax, obj_h = pad_one_dim(ymin, ymax, 0, img_h)\n    \n    if zoom != 0.0:\n        pad_w = obj_w * zoom / 2\n        pad_h = obj_h * zoom / 2\n        xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad_w, img_w)\n        ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad_h, img_h)\n    \n    if try_square:\n        # try pad both sides equaly\n        if obj_w > obj_h:\n            pad = (obj_w - obj_h) / 2\n            ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad, img_h)\n        elif obj_h > obj_w:\n            pad = (obj_h - obj_w) / 2\n            xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad, img_w)\n\n        # if it's still not square, try pad where possible\n        if obj_w > obj_h:\n            pad = obj_w - obj_h\n            ymin, ymax, obj_h = pad_one_dim(ymin, ymax, pad, img_h)\n        elif obj_h > obj_w:\n            pad = obj_h - obj_w\n            xmin, xmax, obj_w = pad_one_dim(xmin, xmax, pad, img_w)\n            \n    return int(xmin), int(ymin), int(xmax), int(ymax)\n\n\n@jit(nopython=True)\ndef bb2wh(bbox):\n    width = bbox[2] - bbox[0]\n    height = bbox[3] - bbox[1]\n    return width, height\n\n\ndef make_x2res(img, bbox):\n    if min(bb2wh(bbox)) < IMG_SIZE_2: return\n    ar = img.width / img.height\n    if ar == 1.0:          \n        tfm_img = resize2x(img)\n    elif 1.0 < ar < 1.15:   \n        tfm_img = center_crop2x(img)\n    elif 1.15 < ar < 1.25:\n        tfm_img = pad_only2x(img)\n    elif 1.25 < ar < 1.5:\n        tfm_img = two_crops2x(img)\n    elif 1.0 < 1/ar < 1.6: \n        tfm_img = top_crop2x(img)\n    else:\n        tfm_img = None\n    return tfm_img\n    \n    \ndef add_sample(samples, label, tfm, imgs, labels):\n    if not samples: return\n    elif isinstance(samples, Image.Image):\n        imgs.append(MyImg(samples, tfm))\n        labels.append(label)\n    elif isinstance(samples, list):\n        imgs.extend([MyImg(s, tfm) for s in samples])\n        labels.extend([label] * len(samples))\n    else: assert False\n\n    \ndef is_valid_file(x):\n    return datasets.folder.has_file_allowed_extension(x, IMG_EXTENSIONS)\n    \n\nclass DogsDataSet(datasets.vision.VisionDataset):\n    def __init__(self, root, transforms, target_transform=None, max_samples=None):\n        super().__init__(root, transform=None)\n        assert isinstance(transforms, list) and len(transforms) == 3\n        self.transforms = transforms\n        self.target_transform = target_transform\n        self.max_samples = max_samples\n        self.classes = {}\n        \n        imgs, labels = self._load_subfolders_images(self.root)   \n        assert len(imgs) == len(labels)\n        if len(imgs) == 0:\n            raise RuntimeError(f'Found 0 files in subfolders of: {self.root}')\n        self.imgs    = imgs\n        self.labels  = labels\n        \n    def _create_or_get_class(self, name):\n        try:\n            label = self.classes[name]\n        except KeyError:\n            label = len(self.classes)\n            self.classes[name] = label\n        return label\n            \n    def _load_subfolders_images(self, root):\n        NO_ZOOM, LIGHT_ZOOM, MEDIUM_ZOOM = 0.0, 0.08, 0.12\n        n_pad, n_center, n_top, n_2crops, n_skip, n_dup, n_noop = 0, 0, 0, 0, 0, 0, 0\n        imgs, labels, paths = [], [], []\n        \n        add_sample_ = partial(add_sample, imgs=imgs, labels=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                paths.append(path)\n        if self.max_samples:\n            paths = paths[:self.max_samples]\n\n        for path in paths:\n            if not is_valid_file(path): continue\n            img = datasets.folder.default_loader(path)\n            annotation_basename = os.path.splitext(os.path.basename(path))[0]\n            annotation_dirname = next(dirname for dirname in os.listdir(ANNOT_PATH) if dirname.startswith(annotation_basename.split('_')[0]))\n            annotation_filename = os.path.join(ANNOT_PATH, annotation_dirname, annotation_basename)\n            tree = ET.parse(annotation_filename)\n            root = tree.getroot()\n            objects = root.findall('object')\n            for o in objects:\n                name = o.find('name').text\n                label = self._create_or_get_class(name)\n                prev_bbox, tfm_imgs = None, None\n                \n                bbox = calc_bbox(o, img_w=img.width, img_h=img.height, zoom=LIGHT_ZOOM)\n                obj_img = img.crop(bbox)\n                add_sample_(make_x2res(obj_img, bbox), label, 2)\n\n                bbox = calc_bbox(o, img_w=img.width, img_h=img.height)                \n                if min(bb2wh(bbox)) < IMG_SIZE:\n                    # don't want pixel mess in gen imgs\n                    n_skip += 1; continue\n                obj_img = img.crop(bbox)\n                ar = obj_img.width / obj_img.height\n                if ar == 1.0:\n                    tfm_imgs = [resize(obj_img)]; n_noop += 1\n                elif 1.0 < ar < 1.3:\n                    tfm_imgs = [center_crop(obj_img), pad_only(obj_img)]\n                    n_center += 1; n_pad += 1\n                elif 1.3 <= ar < 1.5:\n                    tfm_imgs = two_crops(obj_img) + [pad_only(obj_img)]\n                    n_2crops += 2; n_pad += 1\n                elif 1.5 <= ar < 2.0:\n                    tfm_imgs = two_crops(obj_img); n_2crops += 2\n                elif 1.0 < 1/ar < 1.5:\n                    tfm_imgs = [top_crop(obj_img), pad_only(obj_img)]\n                    n_top += 1; n_pad += 1\n                elif 1.5 <= 1/ar < 1.8:\n                    tfm_imgs = [top_crop(obj_img)]; n_top += 1\n                else:\n                    tfm_imgs = None; n_skip += 1\n                add_sample_(tfm_imgs, label, 0)\n                add_sample_(make_x2res(obj_img, bbox), label, 1)\n                prev_bbox = bbox\n\n                bbox = calc_bbox(o, img_w=img.width, img_h=img.height, zoom=MEDIUM_ZOOM, try_square=False)\n                if bbox == prev_bbox:\n                    n_dup += 1; continue\n                if min(bb2wh(bbox)) < IMG_SIZE_2: continue\n                obj_img = img.crop(bbox)\n                ar = obj_img.width / obj_img.height\n                if 1.3 < ar < 1.5:\n                    tfm_imgs = two_crops(obj_img); n_2crops += 2\n                elif 1.05 < 1/ar < 1.6: # maybe tall\n                    tfm_imgs = top_crop(obj_img); n_top += 1\n                else: continue\n                add_sample_(tfm_imgs, label, 0)\n                add_sample_(make_x2res(obj_img, bbox), label, 1)\n                prev_bbox = bbox\n        \n        n_x1, n_x2 = 0, 0\n        for i, img in enumerate(imgs):\n            if img.size == IMG_SIZE: \n                n_x1 +=1\n            else: \n                n_x2 +=1\n                \n        print(f'Found {len(self.classes)} classes\\nLoaded 64x64 {n_x1} images\\n'\n              f'Loaded 128x128 {n_x2} images\\n')\n        print(f'Pad only: {n_pad}\\nCrop center: {n_center}\\n'\n              f'Crop top: {n_top}\\nCrop 2 times: {n_2crops}\\n'\n              f'Take as-is: {n_noop}\\nSkip: {n_skip}\\nSame bbox: {n_dup}')\n        return imgs, labels\n    \n    def __getitem__(self, index):\n        img = self.imgs[index]\n        label = self.labels[index]\n        tfms = self.transforms[img.tfm]\n        img = tfms(image=img.px)['image']          \n        if self.target_transform:\n            label = self.target_transform(label)\n        return img, label\n\n    def __len__(self):\n        return len(self.imgs)\n    \n    \ndef create_runtime_tfms():\n    mean, std = [0.5]*3, [0.5]*3\n    resize_to_64 = A.SmallestMaxSize(IMG_SIZE, interpolation=cv2.INTER_LANCZOS4)\n    out = [A.HorizontalFlip(p=0.5), A.Normalize(mean=mean, std=std), ToTensor()]\n\n    rand_crop = A.Compose([\n        A.SmallestMaxSize(IMG_SIZE + 8, interpolation=cv2.INTER_LANCZOS4),\n        A.RandomCrop(IMG_SIZE, IMG_SIZE)\n    ])\n\n    affine_1 = A.ShiftScaleRotate(\n        shift_limit=0, scale_limit=0.1, rotate_limit=8, \n        interpolation=cv2.INTER_CUBIC, \n        border_mode=cv2.BORDER_REFLECT_101, p=1.0)\n    affine_1 = A.Compose([affine_1, resize_to_64])\n\n    affine_2 = A.ShiftScaleRotate(\n        shift_limit=0.06, scale_limit=(-0.06, 0.18), rotate_limit=6, \n        interpolation=cv2.INTER_CUBIC, \n        border_mode=cv2.BORDER_REFLECT_101, p=1.0)\n    affine_2 = A.Compose([affine_2, resize_to_64])\n\n    tfm_0 = A.Compose(out)\n    tfm_1 = A.Compose([A.OneOrOther(affine_1, rand_crop, p=1.0), *out])\n    tfm_2 = A.Compose([affine_2, *out])\n    return [tfm_0, tfm_1, tfm_2]\n\n\ndef get_data_loaders(data_root=None, batch_size=32, num_workers=2, shuffle=True,\n                     pin_memory=True, drop_last=True):\n    print('Using dataset root location %s' % data_root)\n    # tfms_x1, tfms_x2 = create_runtime_tfms()\n    train_set = DogsDataSet(data_root, create_runtime_tfms())\n    # Prepare loader; the loaders list is for forward compatibility with\n    # using validation / test splits.\n    loaders = []\n    loader_kwargs = {'num_workers': num_workers, 'pin_memory': pin_memory,\n                     'drop_last': drop_last}  # Default, drop last incomplete batch\n    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=shuffle, **loader_kwargs)\n    loaders.append(train_loader)\n    return loaders","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_rng(seed):\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    np.random.seed(seed)\n\n\n# Utility to peg all roots to a base root\n# If a base root folder is provided, peg all other root folders to it.\ndef update_config_roots(config):\n    if config['base_root']:\n        print('Pegging all root folders to base root %s' % config['base_root'])\n        for key in ['data', 'weights', 'logs', 'samples']:\n            config['%s_root' % key] = '%s/%s' % (config['base_root'], key)\n    return config\n\n\ndef prepare_root(config):\n    for key in ['weights_root', 'logs_root', 'samples_root']:\n        if not os.path.exists(config[key]):\n            print('Making directory %s for %s...' % (config[key], key))\n            os.mkdir(config[key])\n\n\n# Function to join strings or ignore them\n# Base string is the string to link \"strings,\" while strings\n# is a list of strings or Nones.\ndef join_strings(base_string, strings):\n    return base_string.join([item for item in strings if item])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# A highly simplified convenience class for sampling from distributions\n# One could also use PyTorch's inbuilt distributions package.\n# Note that this class requires initialization to proceed as\n# x = Distribution(torch.randn(size))\n# x.init_distribution(dist_type, **dist_kwargs)\n# x = x.to(device,dtype)\n# This is partially based on https://discuss.pytorch.org/t/subclassing-torch-tensor/23754/2\nclass Distribution(torch.Tensor):\n    # Init the params of the distribution\n    def init_distribution(self, dist_type, **kwargs):\n        self.dist_type = dist_type\n        self.dist_kwargs = kwargs\n        if self.dist_type == 'normal':\n            self.mean, self.var = kwargs['mean'], kwargs['var']\n        elif self.dist_type == 'categorical':\n            self.num_categories = kwargs['num_categories']\n\n    def sample_(self):\n        if self.dist_type == 'normal':\n            self.normal_(self.mean, self.var)\n        elif self.dist_type == 'categorical':\n            self.random_(0, self.num_categories)\n            # return self.variable\n\n    # Silly hack: overwrite the to() method to wrap the new object\n    # in a distribution as well\n    def to(self, *args, **kwargs):\n        new_obj = Distribution(self)\n        new_obj.init_distribution(self.dist_type, **self.dist_kwargs)\n        new_obj.data = super().to(*args, **kwargs)\n        return new_obj\n\n\ndef prepare_z_y(G_batch_size, dim_z, nclasses, device='cuda', fp16=False, z_var=1.0):\n    z_ = Distribution(torch.randn(G_batch_size, dim_z, requires_grad=False))\n    z_.init_distribution('normal', mean=0, var=z_var)\n    z_ = z_.to(device, torch.float16 if fp16 else torch.float32)\n\n    if fp16:\n        z_ = z_.half()\n\n    y_ = Distribution(torch.zeros(G_batch_size, requires_grad=False))\n    y_.init_distribution('categorical', num_categories=nclasses)\n    y_ = y_.to(device, torch.int64)\n    return z_, y_\n\n\n# Sample function for use with inception metrics\ndef sample_fn(G, z_, y_, config):\n    with torch.no_grad():\n        z_.sample_()\n        y_.sample_()\n        G_z = G(z_, G.shared(y_))\n        return G_z, y_\n\n\ndef toggle_grad(model, on_or_off):\n    for param in model.parameters():\n        param.requires_grad = on_or_off\n        \n        \n# Simple wrapper that applies EMA to a model. COuld be better done in 1.0 using\n# the parameters() and buffers() module functions, but for now this works\n# with state_dicts using .copy_\nclass apply_ema(object):\n    def __init__(self, source, target, decay=0.9999, start_itr=0):\n        self.source = source\n        self.target = target\n        self.decay = decay\n        # Optional parameter indicating what iteration to start the decay at\n        self.start_itr = start_itr\n        # Initialize target's params to be source's\n        self.source_dict = self.source.state_dict()\n        self.target_dict = self.target.state_dict()\n        print('Initializing EMA parameters to be source parameters...')\n        with torch.no_grad():\n            for key in self.source_dict:\n                self.target_dict[key].data.copy_(self.source_dict[key].data)\n                # target_dict[key].data = source_dict[key].data # Doesn't work!\n\n    def update(self, itr=None):\n        # If an iteration counter is provided and itr is less than the start itr,\n        # peg the ema weights to the underlying weights.\n        if itr and itr < self.start_itr:\n            decay = 0.0\n        else:\n            decay = self.decay\n        with torch.no_grad():\n            for key in self.source_dict:\n                self.target_dict[key].data.copy_(self.target_dict[key].data * decay\n                                                 + self.source_dict[key].data * (1 - decay))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def denorm(x):\n    out = (x + 1) / 2\n    return out.clamp_(0, 1)\n\n\ndef plot_imgs(imgs, cols=8, size=4):\n    n = imgs.shape[0]\n    rows = n // cols\n    _, axes = plt.subplots(figsize=(cols * size, rows * size), ncols=cols, nrows=rows)\n    for i, ax in enumerate(axes.flatten()):\n        img = denorm(imgs[i]).numpy().transpose(1, 2, 0)\n        ax.get_xaxis().set_visible(False)\n        ax.get_yaxis().set_visible(False)\n        ax.imshow(img)\n    plt.subplots_adjust(wspace=0.0, hspace=0.0)\n    plt.show()\n    \n    \ndef save_and_sample(G, G_ema, fixed_z, fixed_y, config):\n    which_G = G_ema if config['ema'] and config['use_ema'] else G\n    with torch.no_grad():\n        fixed_Gz = which_G(fixed_z, which_G.shared(fixed_y))\n        sample_images = fixed_Gz.float().cpu()\n        plot_imgs(sample_images, size=2)\n\n\ndef create_train_fn(G, D, GD, z_, y_, ema, state_dict, config):\n    def train(x, y):\n        G.optim.zero_grad()\n        D.optim.zero_grad()\n        # How many chunks to split x and y into?\n        x = torch.split(x, config['batch_size'])\n        y = torch.split(y, config['batch_size'])\n        counter = 0\n\n        # Optionally toggle D and G's \"require_grad\"\n        if config['toggle_grads']:\n            toggle_grad(D, True)\n            toggle_grad(G, False)\n\n        for step_index in range(config['num_D_steps']):\n            # If accumulating gradients, loop multiple times before an optimizer step\n            D.optim.zero_grad()\n            for accumulation_index in range(config['num_D_accumulations']):\n                z_.sample_()\n                y_.sample_()\n                D_fake, D_real = GD(z_[:config['batch_size']], y_[:config['batch_size']],\n                                    x[counter], y[counter], train_G=False,\n                                    split_D=config['split_D'])\n\n                # Compute components of D's loss, average them, and divide by\n                # the number of gradient accumulations\n                D_loss_real, D_loss_fake = loss_hinge_dis(D_fake, D_real)\n                D_loss = (D_loss_real + D_loss_fake) / float(config['num_D_accumulations'])\n                D_loss.backward()\n                counter += 1\n\n            # Optionally apply ortho reg in D\n            if config['D_ortho'] > 0.0:\n                ortho(D, config['D_ortho'])\n\n            D.optim.step()\n\n        # Optionally toggle \"requires_grad\"\n        if config['toggle_grads']:\n            toggle_grad(D, False)\n            toggle_grad(G, True)\n\n        # Zero G's gradients by default before training G, for safety\n        G.optim.zero_grad()\n\n        # If accumulating gradients, loop multiple times\n        for accumulation_index in range(config['num_G_accumulations']):\n            z_.sample_()\n            y_.sample_()\n            D_fake = GD(z_, y_, train_G=True, split_D=config['split_D'])\n            G_loss = loss_hinge_gen(D_fake) / float(config['num_G_accumulations'])\n            G_loss.backward()\n\n        # Optionally apply modified ortho reg in G\n        if config['G_ortho'] > 0.0:\n            # Don't ortho reg shared, it makes no sense. Really we should blacklist any embeddings for this\n            ortho(G, config['G_ortho'], blacklist=[param for param in G.shared.parameters()])\n        G.optim.step()\n\n        # If we have an ema, update it, regardless of if we test with it or not\n        if config['ema']:\n            ema.update(state_dict['itr'])\n\n        out = {'G_loss': float(G_loss.item()),\n               'D_loss_real': float(D_loss_real.item()),\n               'D_loss_fake': float(D_loss_fake.item())}\n        # Return G's loss and the components of D's loss.\n        return out\n\n    return train\n\n\ndef loss_hinge_dis(dis_fake, dis_real):\n    loss_real = torch.mean(F.relu(1. - dis_real))\n    loss_fake = torch.mean(F.relu(1. + dis_fake))\n    return loss_real, loss_fake\n\n\ndef loss_hinge_gen(dis_fake):\n    loss = -torch.mean(dis_fake)\n    return loss\n\n\n# Apply modified ortho reg to a model\n# This function is an optimized version that directly computes the gradient,\n# instead of computing and then differentiating the loss.\ndef ortho(model, strength=1e-4, blacklist=[]):\n    with torch.no_grad():\n        for param in model.parameters():\n            # Only apply this to parameters with at least 2 axes, and not in the blacklist\n            if len(param.shape) < 2 or any([param is item for item in blacklist]):\n                continue\n            w = param.view(param.shape[0], -1)\n            grad = (2 * torch.mm(torch.mm(w, w.t())\n                                 * (1. - torch.eye(w.shape[0], device=w.device)), w))\n            param.grad.data += strength * grad.view(param.shape)\n\n\n# Default ortho reg\n# This function is an optimized version that directly computes the gradient,\n# instead of computing and then differentiating the loss.\ndef default_ortho(model, strength=1e-4, blacklist=[]):\n    with torch.no_grad():\n        for param in model.parameters():\n            # Only apply this to parameters with at least 2 axes & not in blacklist\n            if len(param.shape) < 2 or param in blacklist:\n                continue\n            w = param.view(param.shape[0], -1)\n            grad = (2 * torch.mm(torch.mm(w, w.t())\n                                 - torch.eye(w.shape[0], device=w.device), w))\n            param.grad.data += strength * grad.view(param.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def run(G, D, GD, loaders, state_dict, config, device, G_ema=None, ema=None):\n    # Prepare noise and randomly sampled label arrays\n    # Allow for different batch sizes in G\n    G_batch_size = max(config['G_batch_size'], config['batch_size'])\n    z_, y_ = prepare_z_y(\n        G_batch_size, G.dim_z, config['n_classes'], device=device, fp16=config['G_fp16'])\n    # Prepare a fixed z & y to see individual sample evolution throghout training\n    fixed_z, fixed_y = prepare_z_y(\n        G_batch_size, G.dim_z, config['n_classes'], device=device, fp16=config['G_fp16'])\n    fixed_z.sample_()\n    fixed_y.sample_()\n    # Loaders are loaded, prepare the training function\n    train = create_train_fn(G, D, GD, z_, y_, ema, state_dict, config)\n\n    print('Beginning training at epoch %d...' % state_dict['epoch'])\n    start_time = time.perf_counter()\n    total_iters = config['num_epochs'] * len(loaders[0])\n\n    # Train for specified number of epochs, although we mostly track G iterations.\n    for epoch in range(state_dict['epoch'], config['num_epochs']):\n        for i, (x, y) in enumerate(loaders[0]):\n            # Increment the iteration counter\n            state_dict['itr'] += 1\n            # Make sure G and D are in training mode, just in case they got set to eval\n            # For D, which typically doesn't have BN, this shouldn't matter much.\n            G.train()\n            D.train()\n            if config['ema']:\n                G_ema.train()\n            if config['D_fp16']:\n                x, y = x.to(device).half(), y.to(device)\n            else:\n                x, y = x.to(device), y.to(device)\n            metrics = train(x, y)\n\n            if not (state_dict['itr'] % config['log_interval']):\n                curr_time = time.perf_counter()\n                curr_time_str = datetime.datetime.fromtimestamp(curr_time).strftime('%H:%M:%S')\n                elapsed = str(datetime.timedelta(seconds=(curr_time - start_time)))\n                log = (\n                    \"[{}] [{}] [{} / {}] Ep {}, \".format(curr_time_str, elapsed, state_dict['itr'], total_iters, epoch) + \n                    ', '.join(['%s : %+4.3f' % (key, metrics[key]) for key in metrics])\n                )\n                print(log)\n\n            # Save weights and copies as configured at specified interval\n            if not (state_dict['itr'] % config['save_every']):\n                if config['G_eval_mode']:\n                    print('Switchin G to eval mode...')\n                    G.eval()\n                    # if config['ema']:\n                        # G_ema.eval()\n                save_and_sample(G, G_ema, fixed_z, fixed_y, config)\n                \n            if config['stop_after'] > 0 and int(time.perf_counter() - start_time) > config['stop_after']:\n                print(\"Time limit reached! Stopping training!\") \n                return\n\n        # Increment epoch counter at end of epoch\n        state_dict['epoch'] += 1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def trunc_trick(bs, z_dim, bound=2.0):\n    z = torch.randn(bs, z_dim)\n    while z.min() < -bound or bound < z.max():\n        z = z.where((-bound < z) & (z < bound), torch.randn_like(z))\n    return z\n\n\ndef collect_bn_stats(G, n_samples, config):\n    im_batch_size = config['n_classes']\n    G.train()\n\n    for i_batch in range(0, n_samples, im_batch_size):\n        with torch.no_grad():\n            z = torch.randn(im_batch_size, G.dim_z, device=device)\n            y = torch.arange(im_batch_size).to(device)\n            images = G(z, G.shared(y)).float().cpu()\n\n\ndef generate_images(out_dir, G, n_images, config):\n    im_batch_size = config['n_classes']\n    z_bound = config['trunc_z']\n    if z_bound > 0.0: print(f'Truncating z to (-{z_bound}, {z_bound})')\n\n    for i_batch in range(0, n_images, im_batch_size):\n        with torch.no_grad():\n            if z_bound > 0.0:\n                z = trunc_trick(im_batch_size, G.dim_z, bound=z_bound).to(device)\n            else:\n                z = torch.randn(im_batch_size, G.dim_z, device=device)\n            y = torch.arange(im_batch_size).to(device)\n            images = G(z, G.shared(y)).float().cpu()\n\n        if i_batch + im_batch_size > n_images:\n            n_last_images = n_images - i_batch\n            print(f'Taking only {n_last_images} images from the last batch...')\n            images = images[:n_last_images]\n\n        for i_image, image in enumerate(images):\n            fname = os.path.join(out_dir, f'image_{i_batch+i_image:05d}.png')\n            image = denorm(image)\n            if config['denoise']:\n                image = image * 256\n                image = image.numpy().transpose(1, 2, 0).astype('uint8')\n                image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)\n                image = cv2.fastNlMeansDenoisingColored(\n                    src=image, dst=None, \n                    h=config['denoise_str_lum'], \n                    hColor=config['denoise_str_col'], \n                    templateWindowSize=config['denoise_kernel_size'], \n                    searchWindowSize=config['denoise_search_window']\n                )\n                cv2.imwrite(fname, image)\n            else:\n                torchvision.utils.save_image(image, fname)\n\n\ndef show_saved_samples(n_images=256, cols=8, size=3):\n    rows = math.ceil(n_images / cols)\n    _, axes = plt.subplots(rows, cols, figsize=(cols * size, rows * size))\n    for i, ax in enumerate(axes.flatten()):\n        img = Image.open(f'{OUT_DIR}/image_{i:05d}.png')\n        ax.get_xaxis().set_visible(False)\n        ax.get_yaxis().set_visible(False)\n        ax.imshow(img)\n    plt.subplots_adjust(wspace=0.0, hspace=0.0)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"activation_dict = {\n    'inplace_relu': nn.ReLU(inplace=True),\n    'relu': nn.ReLU(inplace=False),\n    'ir': nn.ReLU(inplace=True),\n}\n\nconfig = {\n    'num_workers': 2,\n    'pin_memory': True,\n    'shuffle': True,\n    'use_multiepoch_sampler': False,\n\n    # Model\n    'model': 'BigGAN',\n    'G_param': 'SN',\n    'D_param': 'SN',\n    'G_ch': 64,\n    'D_ch': 64,\n    'G_depth': 1,\n    'D_depth': 1,\n    'D_wide': True,\n    'G_shared': False,\n    'shared_dim': 0,\n    'dim_z': 128,\n    'z_var': 1.0,\n    'hier': False,\n    'cross_replica': False,\n    'mybn': False,\n    'G_nl': 'relu',\n    'D_nl': 'relu',\n    'G_attn': '64',\n    'D_attn': '64',\n    'norm_style': 'bn',\n\n    ### Model init stuff ###\n    'seed': 0,\n    'G_init': 'ortho',\n    'D_init': 'ortho',\n    'skip_init': False,\n\n    ### Optimizer stuff ###\n    'G_lr': 5e-5,\n    'D_lr': 2e-4,\n    'G_B1': 0.0,\n    'D_B1': 0.0,\n    'G_B2': 0.999,\n    'D_B2': 0.999,\n\n    ### Batch size, parallel, and precision stuff ###\n    'batch_size': 64,\n    'G_batch_size': 0,\n    'num_G_accumulations': 1,\n    'num_D_steps': 2,\n    'num_D_accumulations': 1,\n    'split_D': False,\n    'num_epochs': 100,\n    'parallel': False,\n    'G_fp16': False,\n    'D_fp16': False,\n    'D_mixed_precision': False,\n    'G_mixed_precision': False,\n    'accumulate_stats': False,\n    'num_standing_accumulations': 16,\n\n    ### Bookkeping stuff ###\n    'G_eval_mode': False,\n    'save_every': 2000,\n    'num_save_copies': 2,\n    'num_best_copies': 2,\n    'which_best': 'IS',\n    'no_fid': False,\n    'test_every': 5000,\n    'num_inception_images': 50000,\n    'hashname': False,\n    'base_root': '',\n    'data_root': '../data',\n    'weights_root': '../weights',\n    'logs_root': '../logs',\n    'samples_root': '../samples',\n    'pbar': 'mine',\n    'name_suffix': '',\n    'experiment_name': '',\n    'config_from_name': False,\n\n    ### EMA Stuff ###\n    'ema': False,\n    'ema_decay': 0.9999,\n    'use_ema': False,\n    'ema_start': 0,\n\n    ### Numerical precision and SV stuff ###\n    'adam_eps': 1e-8,\n    'BN_eps': 1e-5,\n    'SN_eps': 1e-8,\n    'num_G_SVs': 1,\n    'num_D_SVs': 1,\n    'num_G_SV_itrs': 1,\n    'num_D_SV_itrs': 1,\n\n    ### Ortho reg stuff ###\n    'G_ortho': 0.0,  # 1e-4 is default for BigGAN\n    'D_ortho': 0.0,\n    'toggle_grads': True,\n\n    ### Which train function ###\n    'which_train_fn': 'GAN',\n\n    ### Resume training stuff\n    'load_weights': '',\n    'resume': False,\n\n    ### Log stuff ###\n    'logstyle': '%3.3e',\n    'log_G_spectra': False,\n    'log_D_spectra': False,\n    'sv_log_interval': 10,\n    'log_interval': 100,\n}\n\nconfig['resolution'] = 64\nconfig['n_classes'] = 120\nconfig['data_root'] = INPUT_PATH","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Current experiment\nconfig['num_epochs'] = 100\nconfig['num_D_steps'] = 1\nconfig['G_lr'] = 1e-4 * 0.85\nconfig['D_lr'] = 4e-4 * 0.85\nconfig['D_B2'] = 0.999\nconfig['G_B2'] = 0.999\nconfig['G_ch'] = 32\nconfig['D_ch'] = 64\nconfig['G_attn'] = '32'\nconfig['D_attn'] = '32'\nconfig['G_nl'] = 'inplace_relu'\nconfig['D_nl'] = 'inplace_relu'\nconfig['G_shared'] = True\nconfig['shared_dim'] = 128\nconfig['dim_z'] = 120\nconfig['hier'] = True\nconfig['SN_eps'] = 1e-8\nconfig['BN_eps'] = 1e-5\nconfig['adam_eps'] = 1e-8\nconfig['G_ortho'] = 0.0\nconfig['G_init'] = 'ortho'\nconfig['D_init'] = 'ortho'\nconfig['G_eval_mode'] = True\nconfig['save_every'] = 5000\nconfig['log_interval'] = 500\nconfig['batch_size'] = 32\n\nconfig['ema'] = True\nconfig['use_ema'] = True\nconfig['ema_decay'] = 0.9999\nconfig['ema_start'] = 50_000\n\nconfig['stop_after'] = 31_500\n\nconfig['trunc_z'] = 0.0\nconfig['denoise'] = True\nconfig['denoise_str_lum'] = 2.65\nconfig['denoise_str_col'] = 1.0\nconfig['denoise_kernel_size'] = 7\nconfig['denoise_search_window'] = 11","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"config = update_config_roots(config)\ndevice = 'cuda'\nseed_rng(1234)\nprepare_root(config)\ntorch.backends.cudnn.benchmark = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n# Prepare data; the Discriminator's batch size is all that needs to be passed\n# to the dataloader, as G doesn't require dataloading.\n# Note that at every loader iteration we pass in enough data to complete\n# a full D iteration (regardless of number of D steps and accumulations)\nD_batch_size = (config['batch_size'] * config['num_D_steps'] * config['num_D_accumulations'])\nloaders = get_data_loaders(\n    data_root=INPUT_PATH,\n    batch_size=D_batch_size,\n    num_workers=config['num_workers'],\n    shuffle=config['shuffle'],\n    pin_memory=config['pin_memory'],\n    drop_last=True\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"config['G_activation'] = activation_dict[config['G_nl']]\nconfig['D_activation'] = activation_dict[config['D_nl']]\n\nG = Generator(**config).to(device)\nD = Discriminator(**config).to(device)\n\n# If using EMA, prepare it\nif config['ema']:\n    print('Preparing EMA for G with decay of {}'.format(config['ema_decay']))\n    G_ema = Generator(**{**config, 'skip_init': True, 'no_optim': True}).to(device)\n    ema = apply_ema(G, G_ema, config['ema_decay'], config['ema_start'])\nelse:\n    G_ema, ema = None, None\n\nGD = G_D(G, D)\nprint('Number of params in G: {} D: {}'.format(\n    *[sum([p.data.nelement() for p in net.parameters()]) for net in [G, D]]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"state_dict = {\n    'itr': 0,\n    'epoch': 0, \n    'config': config\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"run(G, D, GD, loaders, state_dict, config, device, G_ema=G_ema, ema=ema)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nif not os.path.exists(OUT_DIR):\n    os.mkdir(OUT_DIR)\n    \nwhich_G = G_ema if config['ema'] and config['use_ema'] else G\n\nif config['use_ema']:\n    collect_bn_stats(G_ema, 20_000, config)\n\ngc.collect()\n# G.eval()\ngenerate_images(OUT_DIR, which_G, 10_000, config)\nshutil.make_archive('images', 'zip', OUT_DIR)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls -1 $OUT_DIR | wc -l","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_saved_samples()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!rm -rf $OUT_DIR","execution_count":null,"outputs":[]}],"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}