First Commit

Author: Nishantsgithub

AlphaZeroNetwork.py

@@ -0,0 +1,208 @@
+
+import torch
+import torch.nn as nn
+
+class ConvBlock( nn.Module ):
+    """
+    The block consists of a conv layer, batch normalization layer
+    and relu activation.
+    """
+    
+    def __init__( self, input_channels, num_filters ):
+        """
+        Args:
+            input_channels (int) the number of input channels
+            num_filters (int) the number of filters in the conv layer
+        """
+        super().__init__()
+        self.conv1 = nn.Conv2d( input_channels, num_filters, 3, padding=1 )
+        self.bn1 = nn.BatchNorm2d( num_filters )
+        self.relu1 = nn.ReLU()
+
+    def __call__( self, x ):
+        """
+        Args:
+            x (torch.Tensor) the tensor to apply the layers to.
+        """
+        x = self.conv1( x )
+        x = self.bn1( x )
+        x = self.relu1( x )
+
+        return x
+
+class ResidualBlock( nn.Module ):
+    """
+    A residual block.
+    """
+
+    def __init__( self, num_filters ):
+        """
+        Args:
+            num_filters (int) the number of filters in the conv layers. Assumes this is the
+            same as the number of input channels
+        """
+        super().__init__()
+        self.conv1 = nn.Conv2d( num_filters, num_filters, 3,
+                padding=1 )
+        self.bn1 = nn.BatchNorm2d( num_filters )
+        self.relu1 = nn.ReLU()
+        self.conv2 = nn.Conv2d( num_filters, num_filters, 3,
+                padding=1 )
+        self.bn2 = nn.BatchNorm2d( num_filters )
+        self.relu2 = nn.ReLU()
+
+    def __call__( self, x ):
+        """
+        Args:
+            x (torch.Tensor) the tensor to apply the layers to.
+        """
+        residual = x
+
+        x = self.conv1( x )
+        x = self.bn1( x )
+        x = self.relu1( x )
+        
+        x = self.conv2( x )
+        x = self.bn2( x )
+        x += residual
+        x = self.relu2( x )
+
+        return x
+
+class ValueHead( nn.Module ):
+    """
+    nn.Module for the value head
+    """
+
+    def __init__( self, input_channels ):
+        """
+        Args:
+            input_channels (int) the number of input channels
+        """
+        super().__init__()
+        self.conv1 = nn.Conv2d( input_channels, 1, 1 )
+        self.bn1 = nn.BatchNorm2d( 1 )
+        self.relu1 = nn.ReLU()
+        self.fc1 = nn.Linear( 64, 256 )
+        self.relu2 = nn.ReLU()
+        self.fc2 = nn.Linear( 256, 1 )
+        self.tanh1 = nn.Tanh()
+
+    def __call__( self, x ):
+        """
+        Args:
+            x (torch.Tensor) the tensor to apply the layers to.
+        """
+
+        x = self.conv1( x )
+        x = self.bn1( x )
+        x = self.relu1( x )
+        x = x.view( x.shape[0], 64 )
+        x = self.fc1( x )
+        x = self.relu2( x )
+        x = self.fc2( x )
+        x = self.tanh1( x )
+
+        return x
+
+class PolicyHead( nn.Module ):
+    """
+    nn.Module for the policy head
+    """
+
+    def __init__( self, input_channels ):
+        """
+        Args:
+            input_channels (int) the number of input channels
+        """
+        super().__init__()
+        self.conv1 = nn.Conv2d( input_channels, 2, 1 )
+        self.bn1 = nn.BatchNorm2d( 2 )
+        self.relu1 = nn.ReLU()
+        self.fc1 = nn.Linear( 128, 4608 )
+    
+    def __call__( self, x ):
+        """
+        Args:
+            x (torch.Tensor) the tensor to apply the layers to.
+        """
+
+        x = self.conv1( x )
+        x = self.bn1( x )
+        x = self.relu1( x )
+        x = x.view( x.shape[0], 128 )
+        x = self.fc1( x )
+
+        return x
+
+class AlphaZeroNet( nn.Module ):
+    """
+    Neural network with AlphaZero architecture.
+    """
+
+    def __init__(self, num_blocks, num_filters ):
+        """
+        Args:
+            num_blocks (int) the number of residual blocks
+            filters_per_conv (int) the number of filters in each conv layer
+        """
+        super().__init__()
+        #The number of input planes is fixed at 16
+        self.convBlock1 = ConvBlock( 16, num_filters )
+
+        residualBlocks = [ ResidualBlock( num_filters ) for i in range( num_blocks ) ]
+
+        self.residualBlocks = nn.ModuleList( residualBlocks )
+
+        self.valueHead = ValueHead( num_filters )
+
+        self.policyHead = PolicyHead( num_filters )
+
+        self.softmax1 = nn.Softmax( dim=1 )
+
+        self.mseLoss = nn.MSELoss()
+        
+        self.crossEntropyLoss = nn.CrossEntropyLoss()
+
+    def __call__( self, x, valueTarget=None, policyTarget=None, policyMask=None ):
+        """
+        Args:
+            x (torch.Tensor) the input tensor.
+            valueTarget (torch.Tensor) the value target.
+            policyTarget (torch.Tensor) the policy target.
+            policyMask (torch.Tensor) the legal move mask
+        """
+
+        x = self.convBlock1( x )
+
+        for block in self.residualBlocks:
+            x = block( x )
+
+        value = self.valueHead( x )
+
+        policy = self.policyHead( x )
+
+        if self.training:
+            
+            valueLoss = self.mseLoss( value, valueTarget )
+
+            policyTarget = policyTarget.view( policyTarget.shape[0] )
+
+            policyLoss = self.crossEntropyLoss( policy, policyTarget )
+            
+            return valueLoss, policyLoss
+
+        else:
+
+            policyMask = policyMask.view( policyMask.shape[0], -1 )
+
+            policy_exp = torch.exp( policy )
+
+            policy_exp *= policyMask.type( torch.float32 )
+
+            policy_exp_sum = torch.sum( policy_exp, dim=1, keepdim=True )
+            
+            policy_softmax = policy_exp / policy_exp_sum
+
+            return value, policy_softmax
+