{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np\nfrom scipy.signal import convolve2d\nimport tensorflow as tf\nimport keras\nfrom keras import layers\nimport matplotlib.pyplot as plt\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nfrom random import randint, random\nimport sys\nimport gc\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test = True\nmax_delta = 5","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\n    Game of life functions\n\"\"\"\n\ndef get_padded_version_n(X):\n    X_pad = np.zeros((X.shape[0], X.shape[-2] + 2, X.shape[-1] + 2), dtype=X.dtype)\n    X_pad[:, 1:-1,1:-1] += X\n    \n    X_pad[:, 0, 1:-1] = X[:, -1, :]\n    X_pad[:, -1, 1:-1] = X[:, 0, :]\n    \n    X_pad[:, 1:-1, 0] = X[:, :, -1]\n    X_pad[:, 1:-1, -1] = X[:, :, 0]\n    \n    X_pad[:, 0, 0] = X[:, -1, -1]\n    X_pad[:, 0, -1] = X[:, -1, 0]\n    X_pad[:, -1, 0] = X[:, 0, -1]\n    X_pad[:, -1, -1] = X[:, 0, 0]\n    \n    return X_pad\n\ndef nConv2d_sw_3x3(X):\n    X_pad = get_padded_version_n(X)\n    N = np.zeros_like(X_pad)\n    \n    N[:, 1:, 1:] += X_pad[:,:-1,:-1]\n    N[:, 1:, :] += X_pad[:,:-1,:]\n    N[:, 1:, :-1] += X_pad[:,:-1,1:]\n\n    N[:, :, 1:] += X_pad[:,:,:-1]\n    N[:, :, :] += X_pad[:,:,:]\n    N[:, :, :-1] += X_pad[:,:,1:]\n\n    N[:, :-1, 1:] += X_pad[:,1:,:-1]\n    N[:, :-1, :] += X_pad[:,1:,:]\n    N[:, :-1, :-1] += X_pad[:,1:,1:]\n    \n    N = N[:,1:-1,1:-1]\n    \n    return N\n\ndef life_step(X):\n    N =  nConv2d_sw_3x3(X) - X\n    return np.logical_or(N == 3, np.logical_and(X, N==2)).astype(np.uint8)\n\nclass GoL_data:\n    def __init__(self, size, max_delta):\n        self.ydim, self.xdim = size\n        self.grid_size = self.ydim * self.xdim\n        self.max_delta = max_delta\n        \n    def iterate_grids(self, batch_size):\n        while 1:\n            # Generate starting states\n            sample_size = 2 * batch_size\n            probs = np.random.random(sample_size) * 0.98 + 0.01\n            grids = (np.random.random((sample_size, self.ydim, self.xdim)) < np.repeat(probs, self.grid_size).reshape((sample_size, self.ydim, self.xdim))).astype(np.uint8)\n\n            # Warmup steps\n            for i in range(5):\n                grids = life_step(grids)\n\n            # Generate deltas\n            deltas = np.random.randint(1, 6, sample_size)\n\n            # Calculating final states\n            initial_grids = grids\n            targets = np.zeros_like(grids)\n            for i in range(5):\n                grids = life_step(grids)\n                step_mask = deltas == i+1\n                targets[step_mask] = grids[step_mask]\n\n            # Keeping only non-empty states\n            keep_mask = targets.sum(axis=(-2, -1)) > 0\n            initial_grids = initial_grids[keep_mask]\n            targets = targets[keep_mask]\n            deltas = deltas[keep_mask]\n\n            yield initial_grids[:batch_size], targets[:batch_size], deltas[:batch_size]\n            \ndef evaluate_gol(I, F, d):\n    Fh = np.zeros_like(I)\n    s = I\n    for k in range(max_delta):\n        s = life_step(s)\n        sm = d == k+1\n        Fh[sm] = s[sm]\n    acc = (F == Fh).mean()\n    return 1 - acc\n\nif test:\n    gen = GoL_data((25, 25), max_delta)\n    i, f, d = next(iter(gen.iterate_grids(1000)))\n        \n    assert evaluate_gol(i, f, d) == 0.0, \"build_grids is wrong\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\n    GAN functions\n    Forked from https://github.com/eriklindernoren/PyTorch-GAN/blob/master/implementations/gan/gan.py\n\"\"\"\n\nif torch.cuda.is_available():  \n    dev = \"cuda:0\" \nelse:  \n    dev = \"cpu\"\ndevice = torch.device(dev)\n\ndtype = torch.float32\n\nn_steps = 50000\nbatch_size = 128\nlr = 0.0002\nb1 = 0.5\nb2 = 0.999\ndisplay_frequency = 50\nsave_frequency = 100\n\nn_layers = 3\nn_filters = 128\nlatent_dim = 10\n\ndef format_gen_input(F, d):\n    F = np.expand_dims(F, axis=1).astype(np.float32) * 2 - 1\n    d_indic = np.repeat(np.eye(max_delta)[d-1], F.shape[-2] * F.shape[-1]).reshape((F.shape[0], max_delta, F.shape[-2], F.shape[-1])).astype(np.float32) * 2 - 1\n    z = np.concatenate((F, d_indic), axis=1)\n    return torch.tensor(z, dtype=dtype, device=device)\n\ndef format_dis_input(I, d, F):\n    d_indic = np.repeat(np.eye(max_delta)[d-1], F.shape[-2] * F.shape[-1]).reshape((F.shape[0], max_delta, F.shape[-2], F.shape[-1]))\n    if torch.is_tensor(I):\n        return torch.cat((I, torch.tensor(d_indic, dtype=dtype, device=device) * 2 - 1, torch.tensor(np.expand_dims(F, axis=1), dtype=dtype, device=device) * 2 - 1), 1)\n    return torch.tensor(np.concatenate((np.expand_dims(I, axis=1), d_indic, np.expand_dims(F, axis=1)), axis=1), dtype=dtype, device=device) * 2 - 1\n\ndef weights_init(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        nn.init.normal_(m.weight.data, 0.0, 0.02)\n    elif classname.find('BatchNorm') != -1:\n        nn.init.normal_(m.weight.data, 1.0, 0.02)\n        nn.init.constant_(m.bias.data, 0)\n        \ndef count_parameters(model):\n    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n\nclass ConvBlock(nn.Module):\n    def __init__(self, infilters, filters):\n        super(ConvBlock, self).__init__()\n        self.conv = nn.Conv2d(infilters, filters, 3, stride=1, padding=1)\n        self.bn = nn.BatchNorm2d(filters)\n\n    def forward(self, s):\n        return F.relu(self.bn(self.conv(s)))\n\nclass ResBlock(nn.Module):\n    def __init__(self, infilters, filters):\n        super(ResBlock, self).__init__()\n        self.conv1 = nn.Conv2d(infilters, filters, kernel_size=3, stride=1,\n                     padding=1, bias=False)\n        self.bn1 = nn.BatchNorm2d(filters)\n        self.conv2 = nn.Conv2d(filters, filters, kernel_size=3, stride=1,\n                     padding=1, bias=False)\n        self.bn2 = nn.BatchNorm2d(filters)\n\n    def forward(self, x):\n        residual = x\n        out = self.conv1(x)\n        out = F.relu(self.bn1(out))\n        out = self.conv2(out)\n        out = self.bn2(out)\n        out += residual\n        out = F.relu(out)\n        return out\n    \nclass OutBlock(nn.Module):\n    def __init__(self, infilters):\n        super(OutBlock, self).__init__()\n        self.conv = nn.Conv2d(infilters, 1, kernel_size=1)\n    \n    def forward(self,s):\n        return F.tanh(self.conv(s))\n\nclass Generator(nn.Module):\n    def __init__(self, n_res_layers, n_filters):\n        super(Generator, self).__init__()\n        self.n_res_layers = n_res_layers\n        \n        self.convt1 = nn.ConvTranspose2d(latent_dim, n_filters * 4, 3, 1, 0, bias=False)\n        self.bnormt1 = nn.BatchNorm2d(n_filters * 4)\n        \n        self.convt2 = nn.ConvTranspose2d(n_filters * 4, n_filters * 2, 3, 2, 0, bias=False)\n        self.bnormt2 = nn.BatchNorm2d(n_filters * 2)\n        \n        self.convt3 = nn.ConvTranspose2d(n_filters * 2, n_filters, 3, 2, 1, bias=False)\n        self.bnormt3 = nn.BatchNorm2d(n_filters)\n        \n        self.convt4 = nn.ConvTranspose2d(n_filters, int(n_filters/2), 3, 2, 1, bias=False)\n        self.bnormt4 = nn.BatchNorm2d(int(n_filters/2))\n        \n        nn.init.normal_(self.convt1.weight.data, 0.0, 0.02)\n        nn.init.normal_(self.convt2.weight.data, 0.0, 0.02)\n        nn.init.normal_(self.convt3.weight.data, 0.0, 0.02)\n        nn.init.normal_(self.convt4.weight.data, 0.0, 0.02)\n        \n        nn.init.normal_(self.bnormt1.weight.data, 1.0, 0.02)\n        nn.init.constant_(self.bnormt1.bias.data, 0)\n        nn.init.normal_(self.bnormt2.weight.data, 1.0, 0.02)\n        nn.init.constant_(self.bnormt2.bias.data, 0)\n        nn.init.normal_(self.bnormt3.weight.data, 1.0, 0.02)\n        nn.init.constant_(self.bnormt3.bias.data, 0)\n        nn.init.normal_(self.bnormt4.weight.data, 1.0, 0.02)\n        nn.init.constant_(self.bnormt4.bias.data, 0)\n        \n        self.convi = ConvBlock(6, int(n_filters/2))\n        \n        self.blocks = []\n        for _ in range(n_res_layers):\n            self.blocks.append(ResBlock(n_filters, n_filters))\n        self.blocks = nn.ModuleList(self.blocks)\n        self.outblock = OutBlock(n_filters)\n    \n    def forward(self,s):\n        z = torch.randn(s.shape[0], latent_dim, 1, 1, dtype=dtype, device=device)\n        \n        z = self.convt1(z)\n        z = self.bnormt1(z)\n        z = F.relu(z)\n        \n        z = self.convt2(z)\n        z = self.bnormt2(z)\n        z = F.relu(z)\n        \n        z = self.convt3(z)\n        z = self.bnormt3(z)\n        z = F.relu(z)\n        \n        z = self.convt4(z)\n        z = self.bnormt4(z)\n        z = F.relu(z)\n        \n        s = self.convi(s)\n        \n        s = torch.cat((s, z), 1)\n        \n        for i in range(self.n_res_layers):\n            s = self.blocks[i](s)\n        s = self.outblock(s)\n        return s\n\nclass Discriminator(nn.Module):\n    def __init__(self, ndf):\n        super(Discriminator, self).__init__()\n        self.main = nn.Sequential(\n            # input is (nc) x 64 x 64\n            nn.Conv2d(7, ndf, 4, 2, 1, bias=False),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size. (ndf) x 32 x 32\n            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf * 2),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size. (ndf*2) x 16 x 16\n            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),\n            nn.BatchNorm2d(ndf * 4),\n            nn.LeakyReLU(0.2, inplace=True),\n            # state size. (ndf*8) x 4 x 4\n            nn.Conv2d(ndf * 4, 1, 3, 1, 0, bias=False),\n            nn.Sigmoid()\n        )\n\n    def forward(self, input):\n        return self.main(input)\n\n\n# Loss function\nadversarial_loss = torch.nn.BCELoss()\n\n# Initialize generator and discriminator\ngenerator = Generator(n_layers, n_filters).to(device)\n#generator.apply(weights_init)\nprint(\"Parmeters of Generator: {}\".format(count_parameters(generator)))\ndiscriminator = Discriminator(n_filters).to(device)\nprint(\"Parmeters of Discriminator: {}\".format(count_parameters(discriminator)))\ndiscriminator.apply(weights_init)\n\noptimizer_G = torch.optim.Adam(generator.parameters(), lr=lr, betas=(b1, b2))\noptimizer_D = torch.optim.Adam(discriminator.parameters(), lr=lr, betas=(b1, b2))\n    \ndata_gen = iter(GoL_data((25, 25), max_delta).iterate_grids(batch_size))\nd_loss_h, g_loss_h, lb_h = [], [], []\nfor s in range(n_steps):\n    Is, Fs, ds = next(data_gen)\n    \n    # Adversarial ground truths\n    valid = torch.ones(len(Is), dtype=dtype, device=device)\n    fake = torch.zeros(len(Is), dtype=dtype, device=device)\n    \n    ## Generator training\n    optimizer_G.zero_grad()\n\n    # Conditionned noisy input of discriminator\n    z = format_gen_input(Fs, ds)\n\n    # Generate a batch of states\n    gen_states = generator(z)\n    \n    dis_input = format_dis_input(gen_states, ds, Fs)\n\n    # Loss measures generator's ability to fool the discriminator\n    g_loss = adversarial_loss(discriminator(dis_input), valid)\n\n    g_loss.backward()\n    optimizer_G.step()\n\n    ## Discriminator training\n    optimizer_D.zero_grad()\n    \n    #print(format_dis_input(Is, ds, Fs)[0, :, :5, :5])\n    #print(format_dis_input(gen_states.detach(), ds, Fs)[0, :, :5, :5])\n    #raise Exception(\"oui\")\n\n    # Measure discriminator's ability to classify real from generated samples\n    real_loss = adversarial_loss(discriminator(format_dis_input(Is, ds, Fs)), valid)\n    fake_loss = adversarial_loss(discriminator(format_dis_input(gen_states.detach(), ds, Fs)), fake)\n    d_loss = (real_loss + fake_loss) / 2\n\n    d_loss.backward()\n    optimizer_D.step()\n    \n    d_loss_h.append(d_loss.item())\n    g_loss_h.append(g_loss.item())\n    \n    Infered = (gen_states[:, 0].detach().cpu().numpy() > 0.0).astype(np.uint8)\n    lb_h.append(evaluate_gol(Infered, Fs, ds))\n\n    if s % display_frequency == display_frequency-1:\n        dhl = np.array(d_loss_h[-display_frequency:])\n        ghl = np.array(g_loss_h[-display_frequency:])\n        lbhl = np.array(lb_h[-display_frequency:])\n        print(\"Step {}: [D loss: {:.4f}, {:.4f}] [G loss: {:.4f}, {:.4f}] [LB: {:.4f}, {:.4f}]\".format(\n            s + 1,\n            dhl.mean(),\n            dhl.std(),\n            ghl.mean(),\n            ghl.std(),\n            lbhl.mean(),\n            lbhl.std()\n        ))\n\n    if s % save_frequency == save_frequency-1:\n        torch.save(generator, 'generator.pt')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(d_loss_h)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(g_loss_h)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.plot(lb_h)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"np.save(\"generator_loss_history\", np.array(g_loss_h))\nnp.save(\"discriminator_loss_history\", np.array(d_loss_h))\nnp.save(\"lb_history\", np.array(lb_h))","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}