models.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from games import Game
torch.manual_seed(0)

#device=torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")

class ResNet(nn.Module):
    def __init__(self,game:Game,num_res_blocks,num_hidden):
        super(ResNet, self).__init__()
        self.start_block=nn.Sequential(
            nn.Conv2d(3,out_channels=num_hidden,kernel_size=3,padding=1),
            nn.BatchNorm2d(num_hidden),
            nn.ReLU(),
        )

        self.backbone=nn.ModuleList(
        [ResBlock(num_hidden=num_hidden) for i in range(num_res_blocks)]
        )

        self.policyHead=nn.Sequential(
            nn.Conv2d(in_channels=num_hidden,out_channels=32,kernel_size=3,padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(32*game.rows*game.cols,game.action_size)
         )
        
        self.valueHead=nn.Sequential(
            nn.Conv2d(in_channels=num_hidden,out_channels=3,kernel_size=3,padding=1),
            nn.BatchNorm2d(3),
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(3*game.rows*game.cols,1),
            nn.Tanh()
        )

    def forward(self,x):
        x=self.start_block(x)
        for res_block in self.backbone:
            x=res_block(x)
        policy=self.policyHead(x)
        value=self.valueHead(x)
        return policy,value

class ResBlock(nn.Module):
    def __init__(self,num_hidden):
        super(ResBlock, self).__init__()
        self.conv1=nn.Conv2d(num_hidden,num_hidden,kernel_size=3,padding=1)
        self.bn1=nn.BatchNorm2d(num_hidden)
        self.conv2=nn.Conv2d(num_hidden,num_hidden,kernel_size=3,padding=1)
        self.bn2=nn.BatchNorm2d(num_hidden)

    def forward(self,x):
        identity=x
        out=self.conv1(x)
        out=self.bn1(out)
        out=F.relu(out)
        out=self.conv2(out)
        out=self.bn2(out)
        out+=identity
        out=F.relu(out)
        return out