update directory structure

models/__init__.py

@@ -0,0 +1 @@
+from .alphago import *

models/alphago.py

@@ -0,0 +1,52 @@
+import torch
+import torch.nn as nn
+
+class AlphaGoModel(nn.Module):
+    def __init__(self, input_shape, is_policy_net=False,
+                #  num_filters=192, first_kernel_size=5, other_kernel_size=3):
+                num_filters=192, first_kernel_size=5, other_kernel_size=3, num_classes=None):
+        super(AlphaGoModel, self).__init__()
+
+        layers = [
+            nn.Conv2d(input_shape[0], num_filters, first_kernel_size, padding=first_kernel_size//2),
+            nn.ReLU()
+        ]
+
+        for i in range(2, 12):
+            layers.extend([
+                nn.Conv2d(num_filters, num_filters, other_kernel_size, padding=other_kernel_size//2),
+                nn.ReLU()
+            ])
+
+        if is_policy_net:
+            assert num_classes is not None, "num_classes must be provided for policy network"
+            layers.extend([
+                nn.Conv2d(num_filters, 1, kernel_size=1, padding=0),
+                # nn.Softmax(dim=1),
+                nn.Flatten()
+            ])
+        else:
+            layers.extend([
+                nn.Conv2d(num_filters, num_filters, other_kernel_size, padding=other_kernel_size//2),
+                nn.ReLU(),
+                nn.Conv2d(num_filters, 1, kernel_size=1, padding=0),
+                nn.ReLU(),
+                nn.Flatten(),
+                nn.Linear(input_shape[1] * input_shape[2], 256),
+                nn.ReLU(),
+                nn.Linear(256, 1),
+                nn.Tanh()
+            ])
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+# Instantiate policy network
+# alphago_policy_model = AlphaGoModel(input_shape=(3, 19, 19), is_policy_net=True)
+alphago_policy_model = AlphaGoModel(input_shape=(3, 19, 19), is_policy_net=True, num_classes=361)
+
+# Instantiate value network
+alphago_value_model = AlphaGoModel(input_shape=(3, 19, 19), is_policy_net=False)
+

models/alphago_zero.py

@@ -0,0 +1,140 @@
+from keras.layers import *
+from keras.models import Model
+
+
+'''The dual residual architecture is the strongest
+of the architectures tested by DeepMind for AlphaGo
+Zero. It consists of an initial convolutional block,
+followed by a number (40 for the strongest, 20 as 
+baseline) of residual blocks. The network is topped
+off by two "heads", one to predict policies and one
+for value functions.
+'''
+def dual_residual_network(input_shape, blocks=20):
+    inputs = Input(shape=input_shape)
+    first_conv = conv_bn_relu_block(name="init")(inputs)
+    res_tower = residual_tower(blocks=blocks)(first_conv)
+    policy = policy_head()(res_tower)
+    value = value_head()(res_tower)
+    return Model(inputs=inputs, outputs=[policy, value])
+
+
+'''The dual convolutional architecture replaces residual
+blocks from the dual residual architecture with batch-normalized
+convolution layers. The default block size is 12.
+'''
+def dual_conv_network(input_shape, blocks=12):
+    inputs = Input(shape=input_shape)
+    first_conv = conv_bn_relu_block(name="init")(inputs)
+    conv_tower = convolutional_tower(blocks=blocks)(first_conv)
+    policy = policy_head()(conv_tower)
+    value = value_head()(conv_tower)
+    return Model(inputs=inputs, outputs=[policy, value])
+
+
+''' In the separate residual architecture policy and value
+head don't share a common "tail", i.e. there's two sets of
+residual blocks for policy and value networks, respectively. 
+'''
+def separate_residual_network(input_shape, blocks=20):
+    inputs_pol = Input(shape=input_shape)
+    first_conv_pol = conv_bn_relu_block(name="init")(inputs_pol)
+    res_tower_pol = residual_tower(blocks=blocks)(first_conv_pol)
+    policy = policy_head()(res_tower_pol)
+    policy_model = Model(inputs=inputs_pol, outputs=policy)
+
+    inputs_val = Input(shape=input_shape)
+    first_conv_val = conv_bn_relu_block(name="init")(inputs_val)
+    res_tower_val = residual_tower(blocks=blocks)(first_conv_val)
+    value = value_head()(res_tower_val)
+    value_model = Model(inputs=inputs_val, outputs=value)
+    
+    return policy_model, value_model
+
+
+'''The separate convolutional network is structurally identical
+to the separate residual network, except that residual blocks
+are replaced by convolutional blocks.
+'''
+def separate_conv_network(input_shape, blocks=20):
+    inputs_pol = Input(shape=input_shape)
+    first_conv_pol = conv_bn_relu_block(name="init")(inputs_pol)
+    conv_tower_pol =  convolutional_tower(blocks=blocks)(first_conv_pol)
+    policy = policy_head()(conv_tower_pol)
+    policy_model = Model(inputs=inputs_pol, outputs=policy)
+
+    inputs_val = Input(shape=input_shape)
+    first_conv_val = conv_bn_relu_block(name="init")(inputs_val)
+    conv_tower_val = convolutional_tower(blocks=blocks)(first_conv_val)
+    value = value_head()(conv_tower_val)
+    value_model = Model(inputs=inputs_val, outputs=value)
+    
+    return policy_model, value_model
+
+
+def conv_bn_relu_block(name, activation=True, filters=256, kernel_size=(3,3), 
+                       strides=(1,1), padding="same", init="he_normal"):
+    def f(inputs):
+        conv = Conv2D(filters=filters, 
+                      kernel_size=kernel_size,
+                      strides=strides,
+                      padding=padding,
+                      kernel_initializer=init,
+                      data_format='channels_first',
+                      name="{}_conv_block".format(name))(inputs)
+        batch_norm = BatchNormalization(axis=1, name="{}_batch_norm".format(name))(conv)
+        return Activation("relu", name="{}_relu".format(name))(batch_norm) if activation else batch_norm
+    return f    
+
+
+def residual_block(block_num, **args):
+    def f(inputs):
+        res = conv_bn_relu_block(name="residual_1_{}".format(block_num), activation=True, **args)(inputs)
+        res = conv_bn_relu_block(name="residual_2_{}".format(block_num) , activation=False, **args)(res)
+        res = add([inputs, res], name="add_{}".format(block_num))
+        return Activation("relu", name="{}_relu".format(block_num))(res) 
+    return f
+
+
+def residual_tower(blocks, **args):
+    def f(inputs):
+        x = inputs
+        for i in range(blocks):
+            x = residual_block(block_num=i)(x)
+        return x
+    return f
+
+def convolutional_tower(blocks, **args):
+    def f(inputs):
+        x = inputs
+        for i in range(blocks):
+            x = conv_bn_relu_block(name=i)(x)
+        return x
+    return f
+
+
+def policy_head():
+    def f(inputs):
+        conv = Conv2D(filters=2, 
+                      kernel_size=(3, 3),
+                      strides=(1, 1),
+                      padding="same",
+                      name="policy_head_conv_block")(inputs)
+        batch_norm = BatchNormalization(axis=1, name="policy_head_batch_norm")(conv)
+        activation = Activation("relu", name="policy_head_relu")(batch_norm)
+        return Dense(units= 19*19 +1, name="policy_head_dense")(activation)
+    return f    
+
+
+def value_head():
+    def f(inputs):
+        conv = Conv2D(filters=1, 
+                      kernel_size=(1, 1),
+                      strides=(1, 1),
+                      padding="same",
+                      name="value_head_conv_block")(inputs)
+        batch_norm = BatchNormalization(axis=1, name="value_head_batch_norm")(conv)
+        activation = Activation("relu", name="value_head_relu")(batch_norm)
+        dense =  Dense(units= 256, name="value_head_dense", activation="relu")(activation)
+        return Dense(units= 1, name="value_head_output", activation="tanh")(dense)
+    return f