{"cells":[{"metadata":{},"cell_type":"markdown","source":"Architecture was taken from https://github.com/ozanciga/gans-with-pytorch"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport matplotlib.pyplot as plt\n\nimport shutil\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\n\nfrom tqdm import tqdm_notebook as tqdm","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Generator and Discriminator"},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch import nn \nfrom torch.nn.parameter import Parameter\nfrom torch.nn import functional as F\nimport shutil\n\n# TODO for ScaledCrossReplicaBatchNorm2d\nclass _BatchNorm(nn.Module):\n    _version = 2\n\n    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,\n                 track_running_stats=True):\n        super(_BatchNorm, self).__init__()\n        self.num_features = num_features\n        self.eps = eps\n        self.momentum = momentum\n        self.affine = affine\n        self.track_running_stats = track_running_stats\n        if self.affine:\n            self.weight = Parameter(torch.Tensor(num_features))\n            self.bias = Parameter(torch.Tensor(num_features))\n        else:\n            self.register_parameter('weight', None)\n            self.register_parameter('bias', None)\n        if self.track_running_stats:\n            self.register_buffer('running_mean', torch.zeros(num_features))\n            self.register_buffer('running_var', torch.ones(num_features))\n            self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))\n        else:\n            self.register_parameter('running_mean', None)\n            self.register_parameter('running_var', None)\n            self.register_parameter('num_batches_tracked', None)\n        self.reset_parameters()\n\n    def reset_running_stats(self):\n        if self.track_running_stats:\n            self.running_mean.zero_()\n            self.running_var.fill_(1)\n            self.num_batches_tracked.zero_()\n\n    def reset_parameters(self):\n        self.reset_running_stats()\n        if self.affine:\n            self.weight.data.uniform_()\n            self.bias.data.zero_()\n\n    def _check_input_dim(self, input):\n        raise NotImplementedError\n\n    def forward(self, input):\n        self._check_input_dim(input)\n\n        exponential_average_factor = 0.0\n\n        if self.training and self.track_running_stats:\n            self.num_batches_tracked += 1\n            if self.momentum is None:  # use cumulative moving average\n                exponential_average_factor = 1.0 / self.num_batches_tracked.item()\n            else:  # use exponential moving average\n                exponential_average_factor = self.momentum\n\n        return F.batch_norm(\n            input, self.running_mean, self.running_var, self.weight, self.bias,\n            self.training or not self.track_running_stats,\n            exponential_average_factor, self.eps)\n\n    def extra_repr(self):\n        return '{num_features}, eps={eps}, momentum={momentum}, affine={affine}, ' \\\n               'track_running_stats={track_running_stats}'.format(**self.__dict__)\n\n    def _load_from_state_dict(self, state_dict, prefix, metadata, strict,\n                              missing_keys, unexpected_keys, error_msgs):\n        version = metadata.get('version', None)\n\n        if (version is None or version < 2) and self.track_running_stats:\n            # at version 2: added num_batches_tracked buffer\n            #               this should have a default value of 0\n            num_batches_tracked_key = prefix + 'num_batches_tracked'\n            if num_batches_tracked_key not in state_dict:\n                state_dict[num_batches_tracked_key] = torch.tensor(0, dtype=torch.long)\n\n        super(_BatchNorm, self)._load_from_state_dict(\n            state_dict, prefix, metadata, strict,\n            missing_keys, unexpected_keys, error_msgs)\n\n# TODO for ScaledCrossReplicaBatchNorm2d\nclass ScaledCrossReplicaBatchNorm2d(_BatchNorm):\n    r\"\"\"Applies Batch Normalization over a 4D input (a mini-batch of 2D inputs\n    with additional channel dimension) as described in the paper\n    `Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift`_ .\n    .. math::\n        y = \\frac{x - \\mathrm{E}[x]}{ \\sqrt{\\mathrm{Var}[x] + \\epsilon}} * \\gamma + \\beta\n    The mean and standard-deviation are calculated per-dimension over\n    the mini-batches and :math:`\\gamma` and :math:`\\beta` are learnable parameter vectors\n    of size `C` (where `C` is the input size).\n    By default, during training this layer keeps running estimates of its\n    computed mean and variance, which are then used for normalization during\n    evaluation. The running estimates are kept with a default :attr:`momentum`\n    of 0.1.\n    If :attr:`track_running_stats` is set to ``False``, this layer then does not\n    keep running estimates, and batch statistics are instead used during\n    evaluation time as well.\n    .. note::\n        This :attr:`momentum` argument is different from one used in optimizer\n        classes and the conventional notion of momentum. Mathematically, the\n        update rule for running statistics here is\n        :math:`\\hat{x}_\\text{new} = (1 - \\text{momentum}) \\times \\hat{x} + \\text{momemtum} \\times x_t`,\n        where :math:`\\hat{x}` is the estimated statistic and :math:`x_t` is the\n        new observed value.\n    Because the Batch Normalization is done over the `C` dimension, computing statistics\n    on `(N, H, W)` slices, it's common terminology to call this Spatial Batch Normalization.\n    Args:\n        num_features: :math:`C` from an expected input of size\n            :math:`(N, C, H, W)`\n        eps: a value added to the denominator for numerical stability.\n            Default: 1e-5\n        momentum: the value used for the running_mean and running_var\n            computation. Can be set to ``None`` for cumulative moving average\n            (i.e. simple average). Default: 0.1\n        affine: a boolean value that when set to ``True``, this module has\n            learnable affine parameters. Default: ``True``\n        track_running_stats: a boolean value that when set to ``True``, this\n            module tracks the running mean and variance, and when set to ``False``,\n            this module does not track such statistics and always uses batch\n            statistics in both training and eval modes. Default: ``True``\n    Shape:\n        - Input: :math:`(N, C, H, W)`\n        - Output: :math:`(N, C, H, W)` (same shape as input)\n    Examples::\n        >>> # With Learnable Parameters\n        >>> m = nn.BatchNorm2d(100)\n        >>> # Without Learnable Parameters\n        >>> m = nn.BatchNorm2d(100, affine=False)\n        >>> input = torch.randn(20, 100, 35, 45)\n        >>> output = m(input)\n    .. _`Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift`:\n        https://arxiv.org/abs/1502.03167\n    \"\"\"\n\n    def _check_input_dim(self, input):\n        if input.dim() != 4:\n            raise ValueError('expected 4D input (got {}D input)'\n                             .format(input.dim()))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch.optim.optimizer import Optimizer, required\n\nfrom torch.autograd import Variable\nimport torch.nn.functional as F\nfrom torch import nn\nfrom torch import Tensor\nfrom torch.nn import Parameter\n\ndef 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":{"trusted":true},"cell_type":"code","source":"\nimport torch\n\nfrom torch import nn\nfrom torch.nn import init\nfrom torch.nn import functional as F\n\nimport functools\nfrom torch.autograd import Variable\n\n\nclass Spectral_Norm:\n    def __init__(self, name):\n        self.name = name\n\n    def compute_weight(self, module):\n        weight = getattr(module, self.name + '_orig')\n        u = getattr(module, self.name + '_u')\n        size = weight.size()\n        weight_mat = weight.contiguous().view(size[0], -1)\n        with torch.no_grad():\n            v = weight_mat.t() @ u\n            v = v / v.norm()\n            u = weight_mat @ v\n            u = u / u.norm()\n        sigma = u @ weight_mat @ v\n        weight_sn = weight / sigma\n        # weight_sn = weight_sn.view(*size)\n\n        return weight_sn, u\n\n    @staticmethod\n    def apply(module, name):\n        fn = Spectral_Norm(name)\n\n        weight = getattr(module, name)\n        del module._parameters[name]\n        module.register_parameter(name + '_orig', weight)\n        input_size = weight.size(0)\n        u = weight.new_empty(input_size).normal_()\n        module.register_buffer(name, weight)\n        module.register_buffer(name + '_u', u)\n\n        module.register_forward_pre_hook(fn)\n\n        return fn\n\n    def __call__(self, module, input):\n        weight_sn, u = self.compute_weight(module)\n        setattr(module, self.name, weight_sn)\n        setattr(module, self.name + '_u', u)\n\n\ndef spectral_norm(module, name='weight'):\n    Spectral_Norm.apply(module, name)\n\n    return module\n\n\ndef spectral_init(module, gain=1):\n    init.xavier_uniform_(module.weight, gain)\n    if module.bias is not None:\n        module.bias.data.zero_()\n\n    return spectral_norm(module)\n\ndef init_linear(linear):\n    init.xavier_uniform_(linear.weight)\n    linear.bias.data.zero_()\n\n\ndef init_conv(conv, glu=True):\n    init.xavier_uniform_(conv.weight)\n    if conv.bias is not None:\n        conv.bias.data.zero_()\n\n\ndef leaky_relu(input):\n    return F.leaky_relu(input, negative_slope=0.2)\n\nclass SelfAttention(nn.Module):\n    \"\"\" Self attention Layer\"\"\"\n    def __init__(self,in_dim,activation=F.relu):\n        super(SelfAttention,self).__init__()\n        self.chanel_in = in_dim\n        self.activation = activation\n        \n        self.query_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)\n        self.key_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim//8 , kernel_size= 1)\n        self.value_conv = nn.Conv2d(in_channels = in_dim , out_channels = in_dim , kernel_size= 1)\n        self.gamma = nn.Parameter(torch.zeros(1))\n\n        self.softmax  = nn.Softmax(dim=-1) #\n\n        init_conv(self.query_conv)\n        init_conv(self.key_conv)\n        init_conv(self.value_conv)\n        \n    def forward(self,x):\n        \"\"\"\n            inputs :\n                x : input feature maps( B X C X W X H)\n            returns :\n                out : self attention value + input feature \n                attention: B X N X N (N is Width*Height)\n        \"\"\"\n        m_batchsize,C,width ,height = x.size()\n        proj_query  = self.query_conv(x).view(m_batchsize,-1,width*height).permute(0,2,1) # B X CX(N)\n        proj_key =  self.key_conv(x).view(m_batchsize,-1,width*height) # B X C x (*W*H)\n        energy =  torch.bmm(proj_query,proj_key) # transpose check\n        attention = self.softmax(energy) # BX (N) X (N) \n        proj_value = self.value_conv(x).view(m_batchsize,-1,width*height) # B X C X N\n\n        out = torch.bmm(proj_value,attention.permute(0,2,1) )\n        out = out.view(m_batchsize,C,width,height)\n        \n        out = self.gamma*out + x\n        return out\n\n\n\nclass ConditionalNorm(nn.Module):\n    def __init__(self, in_channel, n_condition=148):\n        super().__init__()\n\n        self.bn = nn.BatchNorm2d(in_channel, affine=False)\n\n        self.embed = nn.Linear(n_condition, in_channel* 2)\n        self.embed.weight.data[:, :in_channel] = 1\n        self.embed.weight.data[:, in_channel:] = 0\n\n    def forward(self, input, class_id):\n        out = self.bn(input)\n        # print(class_id.dtype)\n        # print('class_id', class_id.size()) # torch.Size([4, 148])\n        # print(out.size()) #torch.Size([4, 128, 4, 4])\n        # class_id = torch.randn(4,1)\n        # print(self.embed)\n        embed = self.embed(class_id)\n        # print('embed', embed.size())\n        gamma, beta = embed.chunk(2, 1)\n        gamma = gamma.unsqueeze(2).unsqueeze(3)\n        beta = beta.unsqueeze(2).unsqueeze(3)\n        # print(beta.size())\n        out = gamma * out + beta\n\n        return out\n\n\nclass GBlock(nn.Module):\n    def __init__(self, in_channel, out_channel, kernel_size=[3, 3],\n                 padding=1, stride=1, n_class=None, bn=True,\n                 activation=F.relu, upsample=True, downsample=False):\n        super().__init__()\n\n        gain = 2 ** 0.5\n\n        self.conv0 = SpectralNorm(nn.Conv2d(in_channel, out_channel,\n                                             kernel_size, stride, padding,\n                                             bias=True if bn else True))\n        self.conv1 = SpectralNorm(nn.Conv2d(out_channel, out_channel,\n                                             kernel_size, stride, padding,\n                                             bias=True if bn else True))\n\n        self.skip_proj = False\n        if in_channel != out_channel or upsample or downsample:\n            self.conv_sc = SpectralNorm(nn.Conv2d(in_channel, out_channel,\n                                                   1, 1, 0))\n            self.skip_proj = True\n\n        self.upsample = upsample\n        self.downsample = downsample\n        self.activation = activation\n        self.bn = bn\n        if bn:\n            self.HyperBN = ConditionalNorm(in_channel, 148)\n            self.HyperBN_1 = ConditionalNorm(out_channel, 148)\n\n    def forward(self, input, condition=None):\n        out = input\n\n        if self.bn:\n            # print('condition',condition.size()) #condition torch.Size([4, 148])\n            out = self.HyperBN(out, condition)\n        out = self.activation(out)\n        if self.upsample:\n            # TODO different form papers\n            out = F.upsample(out, scale_factor=2)\n        out = self.conv0(out)\n        if self.bn:\n            out = self.HyperBN_1(out, condition)\n        out = self.activation(out)\n        out = self.conv1(out)\n\n        if self.downsample:\n            out = F.avg_pool2d(out, 2)\n\n        if self.skip_proj:\n            skip = input\n            if self.upsample:\n                # TODO different form papers\n                skip = F.upsample(skip, scale_factor=2)\n            skip = self.conv_sc(skip)\n            if self.downsample:\n                skip = F.avg_pool2d(skip, 2)\n\n        else:\n            skip = input\n\n        return out + skip\n\n\nclass Generator(nn.Module):\n    def __init__(self, code_dim=100, n_class=1000, chn=96, debug=False):\n        super().__init__()\n\n        self.linear = SpectralNorm(nn.Linear(n_class, 128, bias=False))\n        \n        if debug:\n            chn = 8\n\n        self.first_view = 16 * chn\n\n        self.G_linear = SpectralNorm(nn.Linear(20, 4 * 4 * 16 * chn))\n\n        self.conv = nn.ModuleList([\n                                #GBlock(16*chn, 16*chn, n_class=n_class),\n                                GBlock(16*chn, 8*chn, n_class=n_class),\n                                GBlock(8*chn, 4*chn, n_class=n_class),\n                                GBlock(4*chn, 2*chn, n_class=n_class),\n                                SelfAttention(2*chn),\n                                GBlock(2*chn, 1*chn, n_class=n_class)])\n\n        # TODO impl ScaledCrossReplicaBatchNorm \n        self.ScaledCrossReplicaBN = ScaledCrossReplicaBatchNorm2d(1*chn)\n        self.colorize = SpectralNorm(nn.Conv2d(1*chn, 3, [3, 3], padding=1))\n\n    def forward(self, input, class_id):\n        codes = torch.split(input, 20, 1)\n        class_emb = self.linear(class_id)  # 128\n\n        out = self.G_linear(codes[0])\n        # out = out.view(-1, 1536, 4, 4)\n        out = out.view(-1, self.first_view, 4, 4)\n        ids = 1\n        for i, conv in enumerate(self.conv):\n            if isinstance(conv, GBlock):\n                \n                conv_code = codes[ids]\n                ids = ids+1\n                condition = torch.cat([conv_code, class_emb], 1)\n                # print('condition',condition.size()) #torch.Size([4, 148])\n                out = conv(out, condition)\n\n            else:\n                out = conv(out)\n\n        out = self.ScaledCrossReplicaBN(out)\n        out = F.relu(out)\n        out = self.colorize(out)\n\n        return F.tanh(out)\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, n_class=1000, chn=96, debug=False):\n        super().__init__()\n\n        def conv(in_channel, out_channel, downsample=True):\n            return GBlock(in_channel, out_channel,\n                          bn=False,\n                          upsample=False, downsample=downsample)\n\n        gain = 2 ** 0.5\n        \n\n        if debug:\n            chn = 8\n        self.debug = debug\n\n        self.pre_conv = nn.Sequential(SpectralNorm(nn.Conv2d(3, 1*chn, 3,padding=1),),\n                                      nn.ReLU(),\n                                      SpectralNorm(nn.Conv2d(1*chn, 1*chn, 3,padding=1),),\n                                      nn.AvgPool2d(2))\n        self.pre_skip = SpectralNorm(nn.Conv2d(3, 1*chn, 1))\n\n        self.conv = nn.Sequential(conv(1*chn, 1*chn, downsample=True),\n                                  SelfAttention(1*chn),\n                                  conv(1*chn, 2*chn, downsample=True),    \n                                  conv(2*chn, 4*chn, downsample=True),\n                                  conv(4*chn, 8*chn, downsample=True),\n                                  conv(8*chn, 16*chn, downsample=True)\n                                  #,conv(16*chn, 16*chn, downsample=False)\n                                 )\n\n        self.linear = SpectralNorm(nn.Linear(16*chn, 1))\n\n        self.embed = nn.Embedding(n_class, 16*chn)\n        self.embed.weight.data.uniform_(-0.1, 0.1)\n        self.embed = spectral_norm(self.embed)\n\n    def forward(self, input, class_id):\n        \n        out = self.pre_conv(input)\n        out = out + self.pre_skip(F.avg_pool2d(input, 2))\n        # print(out.size())\n        out = self.conv(out)\n        out = F.relu(out)\n        out = out.view(out.size(0), out.size(1), -1)\n        out = out.sum(2)\n        out_linear = self.linear(out).squeeze(1)\n        embed = self.embed(class_id)\n\n        prod = (out * embed).sum(1)\n\n        # if self.debug == debug:\n        #     print('class_id',class_id.size())\n        #     print('out_linear',out_linear.size())\n        #     print('embed', embed.size())\n        #     print('prod', prod.size())\n\n        return out_linear + prod\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nimport os\nimport time\nimport torch\nimport datetime\n\nimport torch.nn as nn\nfrom torch.autograd import Variable\nfrom torchvision.utils import save_image\n\nfrom utils import *\n\nclass Trainer(object):\n    def __init__(self, data_loader, config):\n\n # Data loader\n        self.data_loader = data_loader\n\n        # exact model and loss\n        \n        self.model = config['model']\n        self.adv_loss = config['adv_loss']\n\n        # Model hyper-parameters\n        self.imsize = config['imsize']\n        self.g_num = config['g_num']\n        self.z_dim = config['z_dim']\n        self.g_conv_dim = config['g_conv_dim']\n        self.d_conv_dim = config['d_conv_dim']\n        self.parallel = config['parallel']\n        self.gpus = config['gpus']\n\n        self.lambda_gp = config['lambda_gp']\n        self.total_step = config['total_step']\n        self.d_iters = config['d_iters']\n        self.batch_size = config['batch_size']\n        self.num_workers = config['num_workers']\n        self.g_lr = config['g_lr']\n        self.d_lr = config['d_lr']\n        self.lr_decay = config['lr_decay']\n        self.beta1 = config['beta1']\n        self.beta2 = config['beta2']\n        self.pretrained_model = config['pretrained_model']\n\n        self.dataset = config['dataset']\n        self.use_tensorboard = config['use_tensorboard']\n        self.image_path = config['image_path']\n        self.log_path = config['log_path']\n        self.model_save_path = config['model_save_path']\n        self.sample_path = config['sample_path']\n        self.log_step = config['log_step']\n        self.sample_step = config['sample_step']\n        self.model_save_step = config['model_save_step']\n        self.version = config['version']\n\n        self.n_class = config['n_class']\n        self.chn = config['chn']\n        class_id = config['class_id']\n\n        # Path\n        self.log_path = os.path.join(config['log_path'], self.version)\n        self.sample_path = os.path.join(config['sample_path'], self.version)\n        self.model_save_path = os.path.join(config['model_save_path'], self.version)\n\n        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\n        print('build_model...')\n        self.build_model()\n\n        if self.use_tensorboard:\n            self.build_tensorboard()\n\n        # Start with trained model\n        if self.pretrained_model:\n            print('load_pretrained_model...')\n            self.load_pretrained_model()\n\n\n    def label_sampel(self):\n        label = torch.LongTensor(self.batch_size, 1).random_()%self.n_class\n        one_hot= torch.zeros(self.batch_size, self.n_class).scatter_(1, label, 1)\n        return label.squeeze(1).to(self.device), one_hot.to(self.device)       \n\n    def train(self):\n\n        # Data iterator\n        data_iter = iter(self.data_loader)\n        step_per_epoch = len(self.data_loader)\n        model_save_step = int(self.model_save_step * step_per_epoch)\n\n        # Fixed input for debugging\n        fixed_z = tensor2var(torch.randn(self.batch_size, self.z_dim))\n\n        # Start with trained model\n        if self.pretrained_model:\n            start = self.pretrained_model + 1\n        else:\n            start = 0\n\n        # Start time\n        print('Start   ======  training...')\n        start_time = time.time()\n        for step in range(start, self.total_step):\n\n            # ================== Train D ================== #\n            self.D.train()\n            self.G.train()\n\n            try:\n                real_images, real_labels = next(data_iter)\n            except:\n                data_iter = iter(self.data_loader)\n                real_images, real_labels = next(data_iter)\n            \n            #print(real_labels)\n            #print(real_labels.shape)\n            #print(type(real_labels))\n            # Compute loss with real images\n\n            real_labels = real_labels.to(self.device)\n            real_images = real_images.to(self.device)\n\n            d_out_real = self.D(real_images, real_labels)\n            if self.adv_loss == 'wgan-gp':\n                d_loss_real = - torch.mean(d_out_real)\n            elif self.adv_loss == 'hinge':\n                d_loss_real = torch.nn.ReLU()(1.0 - d_out_real).mean()\n\n            # apply Gumbel Softmax\n            z = torch.randn(self.batch_size, self.z_dim).to(self.device)\n\n            z_class, z_class_one_hot = self.label_sampel()\n \n            fake_images = self.G(z, z_class_one_hot)\n            d_out_fake = self.D(fake_images, z_class)\n\n            if self.adv_loss == 'wgan-gp':\n                d_loss_fake = d_out_fake.mean()\n            elif self.adv_loss == 'hinge':\n                d_loss_fake = torch.nn.ReLU()(1.0 + d_out_fake).mean()\n\n\n            # Backward + Optimize\n            d_loss = d_loss_real + d_loss_fake\n            self.reset_grad()\n            d_loss.backward()\n            self.d_optimizer.step()\n\n\n            if self.adv_loss == 'wgan-gp':\n                # Compute gradient penalty\n                alpha = torch.rand(real_images.size(0), 1, 1, 1).to(self.device).expand_as(real_images)\n                interpolated = Variable(alpha * real_images.data + (1 - alpha) * fake_images.data, requires_grad=True)\n                out = self.D(interpolated, z_class)\n\n                grad = torch.autograd.grad(outputs=out,\n                                           inputs=interpolated,\n                                           grad_outputs=torch.ones(out.size()).to(self.device),\n                                           retain_graph=True,\n                                           create_graph=True,\n                                           only_inputs=True)[0]\n\n                grad = grad.view(grad.size(0), -1)\n                grad_l2norm = torch.sqrt(torch.sum(grad ** 2, dim=1))\n                d_loss_gp = torch.mean((grad_l2norm - 1) ** 2)\n\n                # Backward + Optimize\n                d_loss = self.lambda_gp * d_loss_gp\n\n                self.reset_grad()\n                d_loss.backward()\n                self.d_optimizer.step()\n\n            # ================== Train G and gumbel ================== #\n            # Create random noise\n            z = torch.randn(self.batch_size, self.z_dim).to(self.device)\n            z_class, z_class_one_hot = self.label_sampel()\n            \n            fake_images = self.G(z, z_class_one_hot)\n\n            # Compute loss with fake images\n            g_out_fake = self.D(fake_images, z_class)  # batch x n\n            if self.adv_loss == 'wgan-gp':\n                g_loss_fake = - g_out_fake.mean()\n            elif self.adv_loss == 'hinge':\n                g_loss_fake = - g_out_fake.mean()\n\n            self.reset_grad()\n            g_loss_fake.backward()\n            self.g_optimizer.step()\n\n\n            # Print out log info\n            if (step + 1) % self.log_step == 0:\n                elapsed = time.time() - start_time\n                elapsed = str(datetime.timedelta(seconds=elapsed))\n                print(\"Elapsed [{}], G_step [{}/{}], D_step[{}/{}], d_out_real: {:.4f}, d_out_fake: {:.4f}, g_loss_fake: {:.4f}\".\n                      format(elapsed, step + 1, self.total_step, (step + 1),\n                             self.total_step , d_loss_real.item(), d_loss_fake.item(), g_loss_fake.item()))\n                \n                if self.use_tensorboard:\n                    self.writer.add_scalar('data/d_loss_real', d_loss_real.item(),(step + 1))\n                    self.writer.add_scalar('data/d_loss_fake', d_loss_fake.item(),(step + 1))\n                    self.writer.add_scalar('data/d_loss', d_loss.item(), (step + 1))\n\n                    self.writer.add_scalar('data/g_loss_fake', g_loss_fake.item(), (step + 1))\n\n            # Sample images\n            if (step + 1) % self.sample_step == 0:\n                print('Sample images {}_fake.png'.format(step + 1))\n                fake_images= self.G(fixed_z, z_class_one_hot)\n                save_image(denorm(fake_images.data),\n                           os.path.join(self.sample_path, '{}_fake.png'.format(step + 1)))\n                #fixed_z = tensor2var(torch.randn(self.batch_size, self.z_dim))\n                #fake_images2= self.G(fixed_z, z_class_one_hot)\n                #fake_images2=denorm(fake_images2.data)\n                #plt.imshow(fake_images2[0, :, :, :])\n                #shutil.make_archive('learning', 'zip', '../learning/dogs')\n                \n                \n            if (step+1) % model_save_step==0:\n                torch.save(self.G.state_dict(),\n                           os.path.join(self.model_save_path, '{}_G.pth'.format(step + 1)))\n                torch.save(self.D.state_dict(),\n                           os.path.join(self.model_save_path, '{}_D.pth'.format(step + 1)))\n            \n        if not os.path.exists('../output_images'):\n            os.mkdir('../output_images')\n    \n        im_batch_size = self.batch_size\n        n_images=10000\n\n        for i_batch in range(0, n_images, im_batch_size):\n            #print(i_batch)\n            fixed_z = tensor2var(torch.randn(im_batch_size, self.z_dim))\n            fake_images= self.G(fixed_z, z_class_one_hot)\n            fake_images=denorm(fake_images.data)\n            for i_image in range(fake_images.size(0)):\n                #print(i_image)\n                save_image(fake_images[i_image, :, :, :],os.path.join('../output_images', f'image_{i_batch+i_image:05d}.png'))\n     \n        import shutil\n        shutil.make_archive('images', 'zip', '../output_images')\n\n    def build_model(self):\n        # code_dim=100, n_class=1000\n        self.G = Generator(self.z_dim, self.n_class, chn=self.chn).to(self.device)\n        self.D = Discriminator(self.n_class, chn=self.chn).to(self.device)\n        if self.parallel:\n            print('use parallel...')\n            print('gpuids ', self.gpus)\n            gpus = [int(i) for i in self.gpus.split(',')]\n    \n            self.G = nn.DataParallel(self.G, device_ids=gpus)\n            self.D = nn.DataParallel(self.D, device_ids=gpus)\n\n        # self.G.apply(weights_init)\n        # self.D.apply(weights_init)\n\n        # Loss and optimizer\n        # self.g_optimizer = torch.optim.Adam(self.G.parameters(), self.g_lr, [self.beta1, self.beta2])\n        self.g_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, self.G.parameters()), self.g_lr, [self.beta1, self.beta2])\n        self.d_optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, self.D.parameters()), self.d_lr, [self.beta1, self.beta2])\n\n        self.c_loss = torch.nn.CrossEntropyLoss()\n        # print networks\n        print(self.G)\n        print(self.D)\n\n    def build_tensorboard(self):\n        from tensorboardX import SummaryWriter\n        # from logger import Logger\n        # self.logger = Logger(self.log_path)\n        \n        tf_logs_path = os.path.join(self.log_path, 'tf_logs')\n        self.writer = SummaryWriter(log_dir=tf_logs_path)\n\n\n    def load_pretrained_model(self):\n        self.G.load_state_dict(torch.load(os.path.join(\n            self.model_save_path, '{}_G.pth'.format(self.pretrained_model))))\n        self.D.load_state_dict(torch.load(os.path.join(\n            self.model_save_path, '{}_D.pth'.format(self.pretrained_model))))\n        print('loaded trained models (step: {})..!'.format(self.pretrained_model))\n\n    def reset_grad(self):\n        self.d_optimizer.zero_grad()\n        self.g_optimizer.zero_grad()\n\n    def save_sample(self, data_iter):\n        real_images, _ = next(data_iter)\n        save_image(denorm(real_images), os.path.join(self.sample_path, 'real.png'))\n        \n    \n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import argparse\n\ndef str2bool(v):\n    return v.lower() in ('true')\n\n    # Model hyper-parameters\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torchvision.datasets as dsets\nfrom torchvision import transforms\n# required imports\nimport os\nimport xml.etree.ElementTree as ET\nimport torchvision\nfrom tqdm import tqdm_notebook as tqdm\nimport numpy as np\nimport pandas as pd\nimport os\nimport matplotlib.pyplot as plt\n\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\n\nfrom tqdm import tqdm_notebook as tqdm\n\ndef doggo_loader(path):\n    img = torchvision.datasets.folder.default_loader(path) # default loader\n\n        # Get bounding box\n    annotation_basename = os.path.splitext(os.path.basename(path))[0]\n    annotation_dirname = next(dirname for dirname in os.listdir('../input/annotation/Annotation/') if dirname.startswith(annotation_basename.split('_')[0]))\n    annotation_filename = os.path.join('../input/annotation/Annotation', 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        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    bbox = (xmin, ymin, xmax, ymax)\n    \n        # return cropped image\n    return img.crop(bbox)\nrandom_transforms = [transforms.ColorJitter(), transforms.RandomRotation(degrees=0)]\n\nclass Data_Loader():\n    def __init__(self, train, dataset, image_path, image_size, batch_size, shuf=True):\n        self.dataset = dataset\n        self.path = image_path\n        self.imsize = image_size\n        self.batch = batch_size\n        self.shuf = shuf\n        self.train = train\n\n    def transform(self, resize, totensor, normalize, centercrop):\n        options = []\n        if centercrop:\n            options.append(transforms.CenterCrop(160))\n        if resize:\n            options.append(transforms.Resize((self.imsize,self.imsize)))\n        if totensor:\n            options.append(transforms.ToTensor())\n        if normalize:\n            options.append(transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)))\n        transform = transforms.Compose(options)\n        return transform\n\n    def load_lsun(self, classes=['church_outdoor_train','classroom_train']):\n        transforms = self.transform(True, True, True, False)\n        dataset = dsets.LSUN(self.path, classes=classes, transform=transforms)\n        return dataset\n    \n    def load_imagenet(self):\n        transforms = self.transform(True, True, True, True)\n        dataset = dsets.ImageFolder(self.path+'/imagenet', transform=transforms)\n        return dataset\n\n    def load_celeb(self):\n        transforms = self.transform(True, True, True, True)\n        dataset = dsets.ImageFolder(self.path+'/CelebA', transform=transforms)\n        return dataset\n\n\n\n    def load_off(self):\n        #transforms = self.transform(True, True, True, False)\n        dataset = torchvision.datasets.ImageFolder(\n                '../input/all-dogs/',\n                loader=doggo_loader, # THE CUSTOM LOADER\n                transform=torchvision.transforms.Compose([transforms.Resize(64),\n                                            transforms.CenterCrop(64),\n                                            transforms.RandomHorizontalFlip(p=0),\n                                            transforms.RandomApply(random_transforms, p=0),\n                                            transforms.ToTensor()\n                                            ,transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n                                            ]) # some transformations, add your data preprocessing here\n            )\n                                                  \n        return dataset\n\n    def loader(self):\n        if self.dataset == 'lsun':\n            dataset = self.load_lsun()\n        elif self.dataset == 'imagenet':\n            dataset = self.load_imagenet()\n        elif self.dataset == 'celeb':\n            dataset = self.load_celeb()\n        elif self.dataset == 'off':\n            dataset = self.load_off()\n\n        print('dataset',len(dataset))\n        loader = torch.utils.data.DataLoader(dataset=dataset,\n                                              batch_size=self.batch,\n                                              shuffle=self.shuf,\n                                              num_workers=2,\n                                              drop_last=True)\n        return loader\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport torch\nfrom torch.autograd import Variable\nfrom torch.nn import init\n\ndef make_folder(path, version):\n        if not os.path.exists(os.path.join(path, version)):\n            os.makedirs(os.path.join(path, version))\n\n\ndef tensor2var(x, grad=False):\n    if torch.cuda.is_available():\n        x = x.cuda()\n    return Variable(x, requires_grad=grad)\n\ndef var2tensor(x):\n    return x.data.cpu()\n\ndef var2numpy(x):\n    return x.data.cpu().numpy()\n\ndef denorm(x):\n    out = (x + 1) / 2\n    return out.clamp_(0, 1)\n\ndef weights_init(m):\n    classname = m.__class__.__name__\n    # print(classname)\n    if classname.find('Conv2d') != -1:\n        init.xavier_normal_(m.weight.data)\n        init.constant_(m.bias.data, 0.0)\n    elif classname.find('Linear') != -1:\n        init.xavier_normal_(m.weight.data)\n        init.constant_(m.bias.data, 0.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 8\nimage_size = 64\n\nconfig = { 'model':'sagan',\n                        'adv_loss':'hinge',\n                        'imsize':image_size,\n                        'g_num':5,\n                        'chn':24,\n                        'z_dim':120,\n                        'g_conv_dim':32,\n                        'd_conv_dim':32,\n                        'lambda_gp':10,\n                        'version':'dogs',\n                        'n_class':2,    \n                        'class_id':1,\n\n                            # Training setting\n                        'total_step':180000,\n                        'd_iters':5,\n                        'batch_size':batch_size,\n                        'num_workers':12,\n                        'g_lr':0.00001,\n                        'd_lr':0.00004,\n                        'lr_decay':0.95,\n                        'beta1':0,\n                        'beta2':0.9,\n\n                            # using pretrained\n                        'pretrained_model':None,\n\n                            # Misc\n                        'train':True,\n                        'parallel':True,\n                        'gpus':'0',\n                        'dataset':'off',\n                        'use_tensorboard':True,\n\n                            # Path\n                        'image_path':'../input/all-dogs/',\n                        'log_path':'../logs',\n                        'model_save_path':'../models',\n                        'sample_path':'../learning',\n                        'attn_path':'../attn',\n\n                            # Step size\n                        'log_step':2000,\n                        'sample_step':2000,\n                        'model_save_step':20000}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"if not os.path.exists('../learning'):\n    os.mkdir('../learning')\nif not os.path.exists('../learning/dogs'):\n    os.mkdir('../learning/dogs')\nif not os.path.exists('../models'):\n    os.mkdir('../models')\nif not os.path.exists('../models/dogs'):\n    os.mkdir('../models/dogs')\nif not os.path.exists('../attn'):\n    os.mkdir('../attn')\nif not os.path.exists('../attn/dogs'):\n    os.mkdir('../attn/dogs')\nif not os.path.exists('../logs'):\n    os.mkdir('../logs')\nif not os.path.exists('../logs/dogs'):\n    os.mkdir('../logs/dogs')\n\nimport os\nos.listdir('../learning')\nimport shutil\n\ndata_loader = Data_Loader(config['train'], config['dataset'], config['image_path'], config['imsize'],\n                             config['batch_size'], shuf=config['train'])\ntrainer = Trainer(data_loader.loader(), config)\n\ntrainer.train()\n\nshutil.make_archive('learning', 'zip', '../learning/dogs')\nshutil.make_archive('attn', 'zip', '../attn/dogs')\nshutil.make_archive('models', 'zip', '../models/dogs')","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}