drake/chess_dragon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

misc fixes

Browse Source
This commit is contained in:
drake
2026-08-24 23:19:51 -04:00
parent c40a9667bd
commit 33dc15a597
11 changed files with 1742 additions and 1262 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/net/model.rs
+76 -78
View File
@@ -5,14 +5,12 @@ use burn::nn::conv::Conv2dConfig;
use burn::nn::loss::{MseLoss, Reduction};
use burn::nn::{BatchNorm, BatchNormConfig, LinearConfig, PaddingConfig2d};
use burn::tensor::activation::log_softmax;
use burn::tensor::backend::AutodiffBackend;
use burn::tensor::Transaction;
use burn::train::{InferenceStep, ItemLazy, TrainOutput, TrainStep};
use burn::{
nn::{conv::Conv2d, Linear, Relu},
prelude::*,
};
use burn_ndarray::NdArray;
/*
Input planes:
1-6: your pieces (Pawn, Knight, Bishop, Rook, Queen, King)
@@ -39,11 +37,11 @@ pub struct ModelMetadata {
}
#[derive(Module, Debug)]
pub struct ResidualBlock<B: Backend> {
conv1: Conv2d<B>,
bn1: BatchNorm<B>,
conv2: Conv2d<B>,
bn2: BatchNorm<B>,
pub struct ResidualBlock {
conv1: Conv2d,
bn1: BatchNorm,
conv2: Conv2d,
bn2: BatchNorm,
activation: Relu,
}
@@ -52,8 +50,8 @@ pub struct ResidualBlockConfig {
channels: usize,
}
impl<B: Backend> ResidualBlock<B> {
pub fn new(channels: usize, device: &B::Device) -> Self {
impl ResidualBlock {
pub fn new(channels: usize, device: &Device) -> Self {
Self {
conv1: Conv2dConfig::new([channels, channels], [3, 3])
.with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
@@ -66,7 +64,7 @@ impl<B: Backend> ResidualBlock<B> {
activation: Relu::new(),
}
}
pub fn forward(&self, x: Tensor<B, 4>) -> Tensor<B, 4> {
pub fn forward(&self, x: Tensor<4>) -> Tensor<4> {
let residual = x.clone();
let x = self.conv1.forward(x);
@@ -81,24 +79,24 @@ impl<B: Backend> ResidualBlock<B> {
}
#[derive(Module, Debug)]
pub struct ChessModel<B: Backend> {
pub struct ChessModel {
// Initial convolution
conv: Conv2d<B>,
bn: BatchNorm<B>,
conv: Conv2d,
bn: BatchNorm,
// Residual tower
residual_blocks: Vec<ResidualBlock<B>>,
residual_blocks: Vec<ResidualBlock>,
// Policy head
policy_conv: Conv2d<B>,
policy_bn: BatchNorm<B>,
policy_fc: Linear<B>,
policy_conv: Conv2d,
policy_bn: BatchNorm,
policy_conv2: Conv2d,
// Value head
value_conv: Conv2d<B>,
value_bn: BatchNorm<B>,
value_fc1: Linear<B>,
value_fc2: Linear<B>,
value_conv: Conv2d,
value_bn: BatchNorm,
value_fc1: Linear,
value_fc2: Linear,
activation: Relu,
}
@@ -110,11 +108,7 @@ pub struct ChessModelConfig {
}
impl ChessModelConfig {
pub fn init<B: Backend>(
num_blocks: usize,
channels: usize,
device: &B::Device,
) -> ChessModel<B> {
pub fn init(num_blocks: usize, channels: usize, device: &Device) -> ChessModel {
ChessModel {
conv: Conv2dConfig::new([18, channels], [3, 3]) // 18 plane input
.with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
@@ -126,14 +120,14 @@ impl ChessModelConfig {
.collect(),
// Policy head
policy_conv: Conv2dConfig::new([channels, 2], [1, 1]).init(device),
policy_bn: BatchNormConfig::new(2).init(device),
policy_fc: LinearConfig::new(2 * 8 * 8, 8 * 8 * 73).init(device), // 4672 typical chess move space
policy_conv: Conv2dConfig::new([channels, 32], [1, 1]).init(device),
policy_bn: BatchNormConfig::new(32).init(device),
policy_conv2: Conv2dConfig::new([32, 73], [1, 1]).init(device),
// Value head
value_conv: Conv2dConfig::new([channels, 1], [1, 1]).init(device),
value_bn: BatchNormConfig::new(1).init(device),
value_fc1: LinearConfig::new(1 * 8 * 8, 256).init(device),
value_fc1: LinearConfig::new(8 * 8, 256).init(device),
value_fc2: LinearConfig::new(256, 1).init(device),
activation: Relu::new(),
@@ -141,8 +135,8 @@ impl ChessModelConfig {
}
}
impl<B: Backend> ChessModel<B> {
pub fn forward(&self, x: Tensor<B, 4>) -> (Tensor<B, 2>, Tensor<B, 2>) {
impl ChessModel {
pub fn forward(&self, x: Tensor<4>) -> (Tensor<2>, Tensor<2>) {
let mut x = self.conv.forward(x);
x = self.bn.forward(x);
x = self.activation.forward(x);
@@ -157,8 +151,8 @@ impl<B: Backend> ChessModel<B> {
let mut p = self.policy_conv.forward(x.clone());
p = self.policy_bn.forward(p);
p = self.activation.forward(p);
let mut p = p.reshape([batch_size, 2 * 8 * 8]);
p = self.policy_fc.forward(p);
p = self.policy_conv2.forward(p);
let p = p.reshape([batch_size, 73 * 8 * 8]);
// -------- Value Head --------
let mut v = self.value_conv.forward(x);
@@ -173,10 +167,10 @@ impl<B: Backend> ChessModel<B> {
pub fn forward_chess(
&self,
boards: Tensor<B, 4>, // e.g. [batch, channels, 8, 8]
policy_targets: Tensor<B, 2>, // move distribution (make sure is normalized)
value_targets: Tensor<B, 2>, // scalar evaluation
) -> ChessOutput<B> {
boards: Tensor<4>, // e.g. [batch, channels, 8, 8]
policy_targets: Tensor<2>, // move distribution (make sure is normalized)
value_targets: Tensor<2>, // scalar evaluation
) -> ChessOutput {
let (policy_logits, value) = self.forward(boards);
let log_probs = log_softmax(policy_logits.clone(), 1);
@@ -190,39 +184,44 @@ impl<B: Backend> ChessModel<B> {
let value_loss =
MseLoss::new().forward(value.clone(), value_targets.clone(), Reduction::Mean);
let total_loss = policy_loss + 0.5 * value_loss;
let total_loss = policy_loss.clone() + 0.5 * value_loss.clone();
ChessOutput {
policy_logits,
value,
policy_targets,
value_targets,
policy_loss,
value_loss,
loss: total_loss,
}
}
}
pub struct ChessOutput<B: Backend> {
pub policy_logits: Tensor<B, 2>, // [sample, num_moves (4672)]
pub value: Tensor<B, 2>, // [sample, 1 value]
pub policy_targets: Tensor<B, 2>,
pub value_targets: Tensor<B, 2>, // kept 2d for consistency?
pub loss: Tensor<B, 1>, // [sample]
pub struct ChessOutput {
pub policy_logits: Tensor<2>, // [sample, num_moves (4672)]
pub value: Tensor<2>, // [sample, 1 value]
pub policy_targets: Tensor<2>,
pub value_targets: Tensor<2>, // kept 2d for consistency?
pub policy_loss: Tensor<1>,
pub value_loss: Tensor<1>,
pub loss: Tensor<1>, // [sample]
}
impl<B: Backend> ItemLazy for ChessOutput<B> {
type ItemSync = ChessOutput<NdArray>;
fn sync(self) -> Self::ItemSync {
let [policy_logits, value, policy_targets, value_targets, loss] = Transaction::default()
.register(self.policy_logits)
.register(self.value)
.register(self.policy_targets)
.register(self.value_targets)
.register(self.loss)
.execute()
.try_into()
.expect("Correct amount of tensor data");
impl ItemLazy for ChessOutput {
fn sync(self) -> Self {
let [policy_logits, value, policy_targets, value_targets, policy_loss, value_loss, loss] =
Transaction::default()
.register(self.policy_logits)
.register(self.value)
.register(self.policy_targets)
.register(self.value_targets)
.register(self.policy_loss)
.register(self.value_loss)
.register(self.loss)
.execute()
.try_into()
.expect("Correct amount of tensor data");
let device = &Default::default();
@@ -231,27 +230,29 @@ impl<B: Backend> ItemLazy for ChessOutput<B> {
value: Tensor::from_data(value, device),
policy_targets: Tensor::from_data(policy_targets, device),
value_targets: Tensor::from_data(value_targets, device),
policy_loss: Tensor::from_data(policy_loss, device),
value_loss: Tensor::from_data(value_loss, device),
loss: Tensor::from_data(loss, device),
}
}
}
impl<B: AutodiffBackend> TrainStep for ChessModel<B> {
type Input = ChessBatch<B>;
type Output = ChessOutput<B>;
impl TrainStep for ChessModel {
type Input = ChessBatch;
type Output = ChessOutput;
fn step(&self, batch: ChessBatch<B>) -> TrainOutput<ChessOutput<B>> {
fn step(&self, batch: ChessBatch) -> TrainOutput<ChessOutput> {
let item = self.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
TrainOutput::new(self, item.loss.backward(), item)
}
}
impl<B: Backend> InferenceStep for ChessModel<B> {
type Input = ChessBatch<B>;
type Output = ChessOutput<B>;
impl InferenceStep for ChessModel {
type Input = ChessBatch;
type Output = ChessOutput;
fn step(&self, batch: ChessBatch<B>) -> ChessOutput<B> {
fn step(&self, batch: ChessBatch) -> ChessOutput {
self.forward_chess(batch.states, batch.policy_targets, batch.value_targets)
}
}
@@ -287,14 +288,14 @@ impl TrainingSample {
pub struct ChessBatcher {}
#[derive(Clone, Debug)]
pub struct ChessBatch<B: Backend> {
pub states: Tensor<B, 4>,
pub policy_targets: Tensor<B, 2>,
pub value_targets: Tensor<B, 2>,
pub struct ChessBatch {
pub states: Tensor<4>,
pub policy_targets: Tensor<2>,
pub value_targets: Tensor<2>,
}
impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
fn batch(&self, items: Vec<TrainingSample>, device: &B::Device) -> ChessBatch<B> {
impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
fn batch(&self, items: Vec<TrainingSample>, device: &Device) -> ChessBatch {
let state_tensors = items
.iter()
.map(|item| {
@@ -318,18 +319,15 @@ impl<B: Backend> Batcher<B, TrainingSample, ChessBatch<B>> for ChessBatcher {
*p /= sum;
}
}
Tensor::<B, 2>::from_floats(TensorData::new(policy, [1, 4672]), device).unsqueeze()
Tensor::<2>::from_floats(TensorData::new(policy, [1, 4672]), device)
})
.collect::<Vec<_>>();
let value_target_tensors = items
.iter()
.map(|item| {
Tensor::<B, 2>::from_floats(
TensorData::new(vec![item.value_target], [1, 1]),
device,
)
.reshape([1, 1])
Tensor::<2>::from_floats(TensorData::new(vec![item.value_target], [1, 1]), device)
.reshape([1, 1])
})
.collect::<Vec<_>>();