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
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
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