{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Reverse Game of Life - RNN\n\nThe basic idea is that we use a CNN to make a prediction, then feed the prediction back into the network, combined with a forward play layer, to see if it can generate sufficent hidden state to correct the prediction."},{"metadata":{"trusted":true},"cell_type":"code","source":"# Download git repository and copy to local directory\n# !rm -rf /ai-games/\n!git clone https://github.com/JamesMcGuigan/ai-games/ /ai-games/\n# !cd /ai-games/; git checkout ad2f8cc94865f1be6083ca699d4b62b0cc039435\n!cp -rf /ai-games/puzzles/game_of_life/* ./   # copy code to kaggle notebook\n!rm -rf ./neural_networks/models/\n!rm -rf ./neural_networks/hardcoded/logs\n!cd /ai-games/; git log -n1 ","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"from numba import njit, prange\nfrom scipy.signal import convolve2d\nfrom typing import Union, List, Tuple, Dict, Callable\nfrom itertools import chain, product\n\nimport humanize\nimport itertools\nimport math\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport os\nimport pandas as pd\nimport scipy\nimport scipy.sparse\nimport sys\nimport time\nimport skimage\nimport skimage.measure\nimport pydash\n\nimport torch\nimport numpy as np\nfrom utils.util import *\nfrom utils.plot import *\nfrom utils.game import *\nfrom utils.datasets import *\nfrom utils.tuplize import *\nfrom hashmaps.crop import *\nfrom hashmaps.hash_functions import *\nfrom hashmaps.translation_solver import *\nfrom hashmaps.repeating_patterns import *\nfrom constraint_satisfaction.fix_submission import *\n\n\ndevice   = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device(\"cpu\")\n__file__ = './notebook.py'\nnotebook_start = time.perf_counter()\n\n### Don't wrap console output text\nfrom IPython.display import display, HTML\ndisplay(HTML(\"\"\"<style>\ndiv.output_area pre {\n    width: 10000px;\n}\n</style>\"\"\"))","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"from abc import ABCMeta\nfrom typing import TypeVar\n\nimport numpy as np\nimport torch\nimport torch as pt\nimport torch.nn.functional as F\nfrom torch import nn\n\nfrom neural_networks.GameOfLifeBase import GameOfLifeBase\nfrom neural_networks.hardcoded.GameOfLifeHardcodedReLU1_21 import GameOfLifeHardcodedReLU1_21\nfrom neural_networks.modules.ReLUX import ReLU1\n\n# noinspection PyTypeChecker\nT = TypeVar('T', bound='GameOfLifeFeatures')\nclass CircularPadding(nn.Module):\n    in_channels  = 1\n    out_channels = 1\n    def forward(self, x):\n        assert x.dim() == 4\n        x[:,:,0,:] += x[:,:,-1,:]\n        x[:,:,-1,:] = x[:,:,0,:]\n        x[:,:,:,0] += x[:,:,:,-1]\n        x[:,:,:,-1] = x[:,:,:,0]\n        return x\n\n    \nclass CircularWrapping(nn.Module):\n    in_channels  = 1\n    out_channels = 1\n    def forward(self, x):\n        assert x.dim() == 4\n        x[:,:, 1, :]  += x[:,:,-1, :]\n        x[:,:,-2, :]  += x[:,:, 0, :]\n        x[:,:, :, 1]  += x[:,:, :,-1]\n        x[:,:, :,-2]  += x[:,:, :, 0]\n        x = x[:,:,1:-1,1:-1]\n        return x\n\n    \nclass GameOfLifeTransposeOutput(nn.Module):\n    def __init__(self, out_channels=1):\n        super().__init__()\n        self.device       = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device(\"cpu\")\n        self.forward_play = GameOfLifeHardcodedReLU1_21()\n        self.transpose    = nn.ConvTranspose2d(in_channels=140, out_channels=out_channels, bias=False, kernel_size=(3,3), stride=1)\n        self.wrapping     = CircularWrapping()\n        self.softmax      = nn.Softmax(dim=1)\n            \n        self.in_channels  = self.transpose.in_channels   # 84+56 == 140\n        self.out_channels = self.transpose.out_channels            \n            \n        weights  = torch.tensor([ list(map(int,f'{n:09b}')) for n in range(2**9) ]).reshape(1,512,3,3)\n        weights  = weights.to(torch.float32).to(self.device)\n        is_alive = weights[:,:,1,1] == 1.0\n        is_sum   = [ torch.sum( weights, dim=(2,3) ) == float(n) for n in range(10) ]\n        self.transpose.weight.data = weights[ (is_alive & (is_sum[3] | is_sum[4] )) | (~is_alive & is_sum[3]) ]   \n        self.transpose.weight.data = self.transpose.weight.data.reshape(-1,1,3,3)  # Weights must be 4d\n        self.transpose.weight.data = torch.cat([ self.transpose.weight.data for _ in range(out_channels) ], dim=1)\n        self.transpose.weight.data.requires_grad_(False)                           # Weights are hardcoded\n\n        \n    def forward(self, x):\n        shape = x.shape\n        # x = x.reshape(-1, 1, x.shape[2], x.shape[3])     # single channel dataset\n        x = self.softmax(x)\n        x = self.transpose(x)\n        x = self.wrapping(x)\n        # x = x.reshape(shape[0], -1, shape[2], shape[3])  # multi channel dataset\n        return x\n        \n    \nclass GameOfLifeRules(nn.Module):\n    \"\"\" Count the number of neighbouring cells (excluding self)\"\"\"\n\n    def __init__(self):\n        super().__init__()\n        self.device       = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device(\"cpu\")\n        self.forward_play = GameOfLifeHardcodedReLU1_21()\n        self.backwards    = nn.ModuleList([\n            nn.ConvTranspose2d(in_channels=1, out_channels=84, bias=False, kernel_size=(3,3), stride=2),\n            nn.ConvTranspose2d(in_channels=1, out_channels=56, bias=False, kernel_size=(3,3), stride=2)\n        ])\n        self.forwards = nn.Conv2d(in_channels=1, out_channels=11, bias=False, kernel_size=(3,3), stride=2 )\n        self.padding  = CircularPadding()\n        \n        weights  = torch.tensor([ list(map(int,f'{n:09b}')) for n in range(2**9) ]).reshape(1,512,3,3)\n        weights  = weights.to(torch.float32).to(self.device)\n        is_alive = weights[:,:,1,1] == 1.0\n        is_sum   = [ torch.sum( weights, dim=(2,3) ) == float(n) for n in range(10) ]\n        self.backwards[0].weight.data = weights[  is_alive & (is_sum[3] | is_sum[4] )]              # (1,36,3,3)\n        self.backwards[1].weight.data = weights[ ~is_alive & is_sum[3] ]                            # (1,56,3,3)\n        self.forwards.weight.data     = weights[ is_sum[1] | is_sum[9] | (~is_alive & is_sum[8]) ]  # (1,11,3,3)\n        #self.forwards.weight.data     = weights[ is_sum[9] | (~is_alive & is_sum[8]) ]             # (1,2,3,3)\n\n        for layer in [ self.backwards[0], self.backwards[1], self.forwards ]:\n            layer.weight.data = layer.weight.data.reshape(1,-1,3,3)  # Weights must be 4d\n            layer.weight.data.requires_grad_(False)                  # Weights are hardcoded\n        for layer in [ self.forwards ]:\n            layer.weight.data = layer.weight.data.reshape(-1,1,3,3)  # Weights must be 4d\n        self.to(self.device)\n\n\n    def forward(self, x):\n        shape = x.shape\n        x = x.reshape(-1, 1, x.shape[2], x.shape[3])     # single channel dataset\n        x = torch.cat([\n            self.backwards[0](x),\n            self.backwards[1](1-x),\n        ], dim=1)\n        x = self.padding(x)\n        x = x.reshape(-1, 1, x.shape[2], x.shape[3])     # single channel dataset\n        x = self.forwards(x)\n        x = x.reshape(shape[0], -1, shape[2], shape[3])  # multi channel dataset\n        return x\n\n\n\n\nclass GameOfLifeRNNRules(GameOfLifeBase, metaclass=ABCMeta):\n    \"\"\"\n    \"\"\"\n    def __init__(self, grid_size=25, state_size=1):\n        super().__init__()\n        self.grid_size  = grid_size\n        self.state_size = state_size\n\n        in_channels  = (2+self.state_size) * 1540  # *(36+56)*9*2\n        out_channels = self.state_size+1\n        # mid_channels = out_channels * 9\n\n        \n        # self.criterion = FocalLoss()\n        self.criterion    = nn.BCELoss()\n        # self.criterion  = nn.MSELoss()\n        self.activation   = nn.PReLU()\n        # self.activation.weight.data = torch.tensor([-0.5])  # 11N Solution\n        self.relu1        = ReLU1()\n\n        # discriminator must be at top-level for autograd to work, its weights are added to the savefile\n        self.frozen_layers = { 'discriminator', 'features', 'output' }\n        self.discriminator = GameOfLifeHardcodedReLU1_21()\n        self.features      = GameOfLifeRules()\n        self.output        = GameOfLifeTransposeOutput(out_channels=out_channels) \n\n        # TODO: output through nn.ConvTranspose2d(stride=0) with CircularWrapping (and deal with 27x27 edge wrapping)\n        self.cnn_layers = nn.ModuleList([\n            # Input, Prediction, State\n            # nn.Conv2d( in_channels=in_channels,          out_channels=mid_channels, kernel_size=(1,1)),\n            # nn.ConvTranspose2d(in_channels=mid_channels, out_channels=mid_channels, bias=False, kernel_size=(3,3), stride=1),\n            # CircularPadding(),\n            # nn.Conv2d( in_channels=mid_channels,         out_channels=mid_channels,    kernel_size=(1,1)),\n            # nn.Conv2d( in_channels=in_channels,         out_channels=out_channels*9,    kernel_size=(3,3), padding=1, padding_mode='circular'),\n            nn.Conv2d( in_channels=in_channels,         out_channels=in_channels//2,    kernel_size=(1,1)),\n            nn.Conv2d( in_channels=in_channels//2,      out_channels=in_channels//4,    kernel_size=(1,1)),\n            nn.Conv2d( in_channels=in_channels//4,      out_channels=in_channels//8,    kernel_size=(1,1)),\n            nn.Conv2d( in_channels=in_channels//8,      out_channels=self.output.in_channels, kernel_size=(1,1)),\n        ])\n        self.dense_layers = nn.ModuleList([\n            # nn.Linear(in_features=grid_size*grid_size*(3+1), out_features=grid_size*grid_size*2),\n            # nn.Linear(in_features=grid_size*grid_size*(2+1), out_features=grid_size*grid_size*2),\n            # nn.Linear(in_features=grid_size*grid_size*(2+1), out_features=grid_size*grid_size*3)\n        ])\n\n\n    def predict(self, x, max_steps=10, **kwargs):\n        output = super().predict(x, max_steps=max_steps, **kwargs)\n        output = output.reshape((-1, output.shape[1], output.shape[2], output.shape[3]))\n        output = output[:,0,:,:]\n        output = output.squeeze()\n        return output\n\n\n    def forward(self, x, max_steps=5, early_stopping=False):\n        x = self.cast_inputs(x)\n        inputs = x[:,0:1,:,:]\n        if x.shape[1] == 1:\n            x = torch.cat([\n                inputs,                                                                              # Input\n                torch.zeros((x.shape[0], 1,               x.shape[2], x.shape[3])).to(self.device),  # Prediction\n                torch.zeros((x.shape[0], self.state_size, x.shape[2], x.shape[3])).to(self.device),  # State\n            ], dim=1)\n\n        for n in range(max_steps):\n            x[:,0:2,:,:] = self.relu1(x[:,0:2,:,:])  # ReLU1 on input and prediction\n            x = self.features(x)\n            for layer in self.cnn_layers:\n                # if x.shape[1] == layer.in_channels // 22:\n                #     x = self.features(x)\n                # if x.shape[1] == layer.in_channels - 1:\n                #     x = torch.cat([ inputs, x ], dim=1)\n\n                x = layer(x)\n                x = self.activation(x)  # V shaped - SVM filter lines\n                # x = self.relu1(x)     # Z shaped - cast back to boolean\n\n            # x = self.features(x)\n            shape = x.shape\n            x = x.flatten(1)\n            for layer in self.dense_layers:\n                x = torch.cat([ x, inputs.flatten(1) ], dim=1)\n                x = layer(x)\n                x = self.activation(x)\n            x = x.reshape((shape[0], -1, shape[2], shape[3]))\n            x = self.output(x)  # Transpose prediction  \n            \n            x = torch.cat([ inputs, x ], dim=1)      # append input to output\n            x[:,0:2,:,:] = self.relu1(x[:,0:2,:,:])  # ReLU1 on input and prediction\n\n            if early_stopping and n != max_steps-1:\n                reinputs, prediction, state = x[:,0:1,:,:], x[:,1:2,:,:], x[:,2:,:,:]\n                forwards   = self.discriminator(prediction)\n                if torch.all(torch.eq( inputs, forwards )):\n                    break\n\n        return x\n\n\n    def cell_count_loss(self, boards1, boards2):\n        \"\"\" Return average difference in cell count (per board) squared \"\"\"\n        return torch.mean((torch.mean(boards1.flatten(1), dim=0) - torch.mean(boards2.flatten(1), dim=0)) ** 2)\n\n    def binary_loss(self, x):\n        \"\"\" loss == 0 if all values are either 0 or 1; max loss == 0.5 if all values are 0.5 \"\"\"\n        return 0.5**2 - torch.mean( ( x - 0.5 ) ** 2 )\n        # return 0.5    - torch.mean( torch.abs( x - 0.5 ) )\n\n\n    def loss(self, outputs, expected, inputs, max_steps=5):\n        \"\"\"\n        GameOfLifeReverseOneGAN() computes the backwards timestep\n        discriminator GameOfLifeHardcodedReLU1_21() replays the board again forwards\n        forward_loss is the MSE difference between the backwards prediction and forward play\n        classic_loss biases the network towards the exact solution, but reduces to zero as forward_loss approaches zero\n        sum_loss is a heuristic to guide solution towards the correct cell count and avoid all 0 or all 1 solutions\n        binary_loss penalizes non-binary output\n        \"\"\"\n        losses = torch.zeros((max_steps,), dtype=torch.float, requires_grad=True).to(self.device)\n        for t in range(max_steps):\n            if t != 0:\n                outputs = torch.cat([ inputs, outputs[:,1:,:,:] ], dim=1)  # reset inputs\n                outputs = self(outputs, max_steps=1, early_stopping=False)\n\n            reinputs, prediction, state = outputs[:,0:1,:,:], outputs[:,1:2,:,:], outputs[:,2:,:,:]\n            forwards = self.discriminator(torch.round(prediction))\n\n            forward_loss  = self.criterion(forwards, inputs)   # loss==0 if forward play matches input\n            identity_loss = self.criterion(reinputs, inputs)   # loss==0 if forward play matches input\n\n            cell_count_loss_forwards   = self.cell_count_loss(forwards,   inputs)\n            cell_count_loss_reinputs   = self.cell_count_loss(reinputs,   inputs)\n            cell_count_loss_prediction = self.cell_count_loss(prediction, expected)\n            cell_count_loss = (cell_count_loss_forwards + cell_count_loss_reinputs + cell_count_loss_prediction) / 3\n\n            binary_loss_prediction = self.binary_loss(prediction)\n            binary_loss_reinputs   = self.binary_loss(reinputs)\n            binary_loss = (binary_loss_prediction + binary_loss_reinputs) / 2\n\n            dataset_loss = self.criterion(prediction, expected)       # loss==0 if output matches dataset\n            dataset_loss = F.relu( dataset_loss * (torch.tanh(forward_loss) - 0.01) )  # fade out classic loss\n\n            losses[t] = forward_loss # + identity_loss + ( cell_count_loss + binary_loss + dataset_loss ) / 3\n            losses[t] = losses[t] * (t+1)  # end of sequence outputs are most important\n\n        loss = torch.sum( losses ) / np.sum(np.arange(1,max_steps+1))\n        return loss\n\n\n    # noinspection PyTypeChecker\n    def accuracy(self, outputs, expected, inputs):\n        \"\"\" Accuracy here is based upon if the output matches the input after forward play \"\"\"\n        # return super().accuracy(outputs, expected, inputs)\n        reinputs, prediction, state = outputs[:,0:1,:,:], outputs[:,1:2,:,:], outputs[:,2:,:,:]\n        forwards = self.discriminator(self.cast_int(prediction))\n        return pt.sum( self.cast_bool(forwards) == self.cast_bool(inputs) ).cpu().numpy() / np.prod(outputs.shape)\n\n\n    def unfreeze(self: T) -> T:\n        if not self.loaded: self.load()\n        for name, parameter in self.named_parameters():\n            if not set( name.split('.') ) & set(self.frozen_layers):\n                parameter.requires_grad = True\n            else:\n                parameter.requires_grad = False\n        return self\n\n\n# if __name__ == '__main__':\n#     from neural_networks.train import train\n#     model = GameOfLifeRNNRules(grid_size=25)\n#     train(model, batch_size=25, grid_size=25, reverse_input_output=True)\n#     # model.print_params()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = GameOfLifeRNNRules()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from neural_networks.train import train\n\ntimeout = 60 * 60 - ( time.perf_counter() - notebook_start )\ntrain(model, batch_size=9, grid_size=25, reverse_input_output=True, timeout=timeout)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from neural_networks.train import train\n\ntimeout = 110 * 60 - ( time.perf_counter() - notebook_start )\ntrain(model, batch_size=25, grid_size=25, reverse_input_output=True, timeout=timeout)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.print_params()","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}