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

training fixes

Browse Source
This commit is contained in:
drake
2026-09-09 15:47:57 -04:00
parent d901c9a85d
commit 01192b9fb8
3 changed files with 98 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/main.rs
+4 -2
View File
@@ -16,7 +16,7 @@ use engine::train::{train, TrainingConfig};
fn main() {
let device = Device::cuda(DeviceIndex::default());
let mcts_config = MctsConfig::new(200, 1.0, 0.05, 0.25, 128);
let mcts_config = MctsConfig::new(400, 1.0, 0.05, 0.25, 128);
let adam_config = AdamConfig::new();
@@ -29,8 +29,10 @@ fn main() {
hidden_channels: 48,
num_blocks: 4,
batch_size: 1024, // num positions sampled to a batch (2048)
minibatch_size: 256,
num_train_steps: 10,
num_episodes: 96, // num games generated per iteration (5000)
buffer_max_size: 25_000, // max number of samples in the buffer
buffer_max_size: 200_000, // max number of samples in the buffer
mcts_config,
optimizer: adam_config,
lr: 2e-4,
engine/src/train.rs
+74 -20
View File
@@ -33,6 +33,8 @@ pub struct TrainingConfig {
pub hidden_channels: usize,
pub num_blocks: usize,
pub batch_size: usize, // num positions sampled to a batch (2048)
pub minibatch_size: usize, // num positions sampled to a batch (2048)
pub num_train_steps: usize, // number of training steps of size minibatch_size
pub num_episodes: usize, // num games generated per iteration (5000)
pub buffer_max_size: usize, // max number of samples in the buffer
pub mcts_config: MctsConfig,
@@ -241,49 +243,101 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
let num_samples = std::cmp::min(training_config.batch_size, replay_buffer.len());
println!(
"Finished self-play! Training on {} samples (buffer size: {})...",
"Finished self-play! Training on {} samples {} times with a minibatch size of {} (buffer size: {})...",
num_samples,
training_config.num_train_steps,
training_config.minibatch_size,
replay_buffer.len()
);
let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
indices.shuffle(&mut rng);
let samples = indices
.into_iter()
.take(num_samples)
.map(|i| replay_buffer[i].clone())
.collect::<Vec<TrainingSample>>();
let batcher = ChessBatcher {};
// Instead of one step, do multiple epochs over the sampled batch
for step in 0..training_config.num_train_steps {
// Re-shuffle and create mini-batches within the sampled data
let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
indices.shuffle(&mut rng);
indices.truncate(num_samples);
let mut loss_val: f32 = 0.0;
let mut p_loss_val: f32 = 0.0;
let mut v_loss_val: f32 = 0.0;
// Process in mini-batches of e.g. 256
for chunk in indices.chunks(training_config.minibatch_size) {
let samples: Vec<TrainingSample> =
chunk.iter().map(|&i| replay_buffer[i].clone()).collect();
let batch = batcher.batch(samples, &autodiff_device);
let output =
model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
let grads = output.loss.backward();
let grads = GradientsParams::from_grads(grads, &model);
model = optim.step(training_config.lr, model, grads);
let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
let loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
let p_loss_val = output
if step == training_config.num_train_steps - 1 {
loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
p_loss_val = output
.policy_loss
.clone()
.into_data()
.to_vec::<f32>()
.unwrap()[0];
let v_loss_val = output
v_loss_val = output
.value_loss
.clone()
.into_data()
.to_vec::<f32>()
.unwrap()[0];
}
}
if step == training_config.num_train_steps - 1 {
println!(
"Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
"Iteration {}, last minibatch training loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
iter, loss_val, p_loss_val, v_loss_val
);
println!("Iteration took {}s", iter_start.elapsed().as_secs());
}
}
let grads = output.loss.backward();
let grads = GradientsParams::from_grads(grads, &model);
model = optim.step(training_config.lr, model, grads);
// let mut indices: Vec<usize> = (0..replay_buffer.len()).collect();
// indices.shuffle(&mut rng);
//
// let samples = indices
// .into_iter()
// .take(num_samples)
// .map(|i| replay_buffer[i].clone())
// .collect::<Vec<TrainingSample>>();
//
// let batcher = ChessBatcher {};
//
// let batch = batcher.batch(samples, &autodiff_device);
//
// let output = model.forward_chess(batch.states, batch.policy_targets, batch.value_targets);
//
// let loss_val = output.loss.clone().into_data().to_vec::<f32>().unwrap()[0];
// let p_loss_val = output
// .policy_loss
// .clone()
// .into_data()
// .to_vec::<f32>()
// .unwrap()[0];
// let v_loss_val = output
// .value_loss
// .clone()
// .into_data()
// .to_vec::<f32>()
// .unwrap()[0];
//
// println!(
// "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
// iter, loss_val, p_loss_val, v_loss_val
// );
// println!("Iteration took {}s", iter_start.elapsed().as_secs());
//
// let grads = output.loss.backward();
// let grads = GradientsParams::from_grads(grads, &model);
//
// model = optim.step(training_config.lr, model, grads);
iter += 1;
engine/src/worker.rs
+5 -2
View File
@@ -84,7 +84,8 @@ impl Worker {
// println!("{}", before.elapsed().as_millis());
episode_buffer.push(results);
let temp = if ply_count < 30 { 1.0 } else { 0.2 };
let temp = (1.0 - (ply_count as f32 / 200.0).min(1.0)).max(0.3);
// let temp = if ply_count < 30 { 1.0 } else { 0.2 };
let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
@@ -125,9 +126,11 @@ impl Worker {
_ => 0.0,
};
let blended_target = 0.75 * outcome + 0.25 * result.value;
samples.push(TrainingSample::from_mcts_with_outcome(
result.clone(),
outcome,
blended_target,
));
}