Example #1
0
 def __init__(self, input_channel_num, block_num, channel_num,
              policy_channel_num, board_size):
     super(CategoricalNetwork, self).__init__()
     self.encoder_ = Encoder(input_channel_num, block_num, channel_num)
     self.policy_head_ = PolicyHead(channel_num, policy_channel_num)
     self.value_head_ = ValueHead(channel_num, 51)
     self.encoder_head = EncodeHead(channel_num, channel_num, channel_num)
Example #2
0
    def __init__(self, input_channel_num, block_num, channel_num,
                 policy_channel_num, board_size):
        super().__init__()

        self.board_size = board_size
        self.num_patch = board_size * board_size
        self.first_conv = nn.Conv2d(input_channel_num, channel_num, 1, 1)

        self.mixer_blocks = nn.ModuleList([])
        for _ in range(block_num):
            self.mixer_blocks.append(MixerBlock(channel_num, self.num_patch))
        self.layer_norm = nn.LayerNorm(channel_num)
        self.dim = channel_num
        self.policy_head_ = PolicyHead(channel_num, policy_channel_num)
        self.value_head_ = ValueHead(channel_num, 51)
Example #3
0
 def __init__(self, input_channel_num, block_num, channel_num,
              policy_channel_num, board_size):
     super(PoolFormerModel, self).__init__()
     self.first_conv1x1_ = nn.Conv2d(input_channel_num, channel_num, 1)
     self.blocks_ = basic_blocks(channel_num, block_num)
     self.policy_head_ = PolicyHead(channel_num, policy_channel_num)
     self.value_head_ = ValueHead(channel_num, 51)
Example #4
0
 def __init__(self, input_channel_num, block_num, channel_num, policy_channel_num, board_size):
     super(TransformerModel, self).__init__()
     self.first_encoding_ = torch.nn.Linear(input_channel_num, channel_num)
     encoder_layer = torch.nn.TransformerEncoderLayer(channel_num, nhead=8, dim_feedforward=channel_num * 4, norm_first=True, activation="gelu")
     self.encoder_ = torch.nn.TransformerEncoder(encoder_layer, block_num)
     self.board_size = board_size
     square_num = board_size ** 2
     self.policy_head_ = PolicyHead(channel_num, policy_channel_num)
     self.value_head_ = ValueHead(channel_num, 51)
     self.positional_encoding_ = torch.nn.Parameter(torch.zeros([square_num, 1, channel_num]), requires_grad=True)
Example #5
0
class ScalarNetwork(nn.Module):
    def __init__(self, input_channel_num, block_num, channel_num,
                 policy_channel_num, board_size):
        super(ScalarNetwork, self).__init__()
        self.encoder_ = Encoder(input_channel_num, block_num, channel_num)
        self.policy_head_ = PolicyHead(channel_num, policy_channel_num)
        self.value_head_ = ValueHead(channel_num, 1)

    def forward(self, x):
        x = self.encoder_.forward(x)
        policy = self.policy_head_.forward(x)
        value = self.value_head_.forward(x)
        value = torch.tanh(value)
        return policy, value
Example #6
0
class MLPMixer(nn.Module):
    def __init__(self, input_channel_num, block_num, channel_num,
                 policy_channel_num, board_size):
        super().__init__()

        self.board_size = board_size
        self.num_patch = board_size * board_size
        self.first_conv = nn.Conv2d(input_channel_num, channel_num, 1, 1)

        self.mixer_blocks = nn.ModuleList([])
        for _ in range(block_num):
            self.mixer_blocks.append(MixerBlock(channel_num, self.num_patch))
        self.layer_norm = nn.LayerNorm(channel_num)
        self.dim = channel_num
        self.policy_head_ = PolicyHead(channel_num, policy_channel_num)
        self.value_head_ = ValueHead(channel_num, 51)

    def forward(self, x):
        batch_size = x.shape[0]
        x = self.first_conv.forward(x)
        x = x.permute([0, 2, 3, 1])
        x = x.view([-1, self.board_size * self.board_size, self.dim])
        for mixer_block in self.mixer_blocks:
            x = mixer_block(x)
        x = self.layer_norm(x)
        x = x.permute([0, 2, 1])
        x = x.contiguous()
        x = x.view([batch_size, self.dim, self.board_size, self.board_size])
        policy = self.policy_head_.forward(x)
        value = self.value_head_.forward(x)
        return policy, value

    @torch.jit.export
    def getRepresentations(self, x):
        x = self.first_conv(x)
        x = x.permute([0, 2, 3, 1])
        x = x.view([-1, self.board_size * self.board_size, self.dim])
        result = []
        for mixer_block in self.mixer_blocks:
            x = mixer_block(x)
            result.append(x)
        return result
Example #7
0
class CategoricalNetwork(nn.Module):
    def __init__(self, input_channel_num, block_num, channel_num,
                 policy_channel_num, board_size):
        super(CategoricalNetwork, self).__init__()
        self.encoder_ = Encoder(input_channel_num, block_num, channel_num)
        self.policy_head_ = PolicyHead(channel_num, policy_channel_num)
        self.value_head_ = ValueHead(channel_num, 51)
        self.encoder_head = EncodeHead(channel_num, channel_num, channel_num)

    def forward(self, x):
        x = self.encoder_.forward(x)
        policy = self.policy_head_.forward(x)
        value = self.value_head_.forward(x)
        return policy, value

    @torch.jit.export
    def encode(self, x):
        x = self.encoder_.forward(x)
        x = self.encoder_head.forward(x)
        return x